| 1 | import functools |
| 2 | import inspect |
| 3 | from pyinstrument import Profiler |
| 4 | |
| 5 | def trace_performance(*, show_all=False, color=True, unicode=True): |
| 6 | """ |
| 7 | Decorator that profiles a function and prints a call tree when it finishes. |
| 8 | |
| 9 | Works with both synchronous and asynchronous functions. |
| 10 | """ |
| 11 | |
| 12 | def decorator(func): |
| 13 | is_coro = inspect.iscoroutinefunction(func) |
| 14 | |
| 15 | @functools.wraps(func) |
| 16 | async def async_wrapper(*args, **kwargs): |
| 17 | profiler = Profiler() |
| 18 | profiler.start() |
| 19 | try: |
| 20 | return await func(*args, **kwargs) |
| 21 | finally: |
| 22 | profiler.stop() |
| 23 | print(f"\n=== Performance trace: {func.__module__}.{func.__qualname__} (async) ===") |
| 24 | print( |
| 25 | profiler.output_text( |
| 26 | color=color, |
| 27 | unicode=unicode, |
| 28 | show_all=show_all, |
| 29 | ) |
| 30 | ) |
| 31 | |
| 32 | @functools.wraps(func) |
| 33 | def sync_wrapper(*args, **kwargs): |
| 34 | profiler = Profiler() |
| 35 | profiler.start() |
| 36 | try: |
| 37 | return func(*args, **kwargs) |
| 38 | finally: |
| 39 | profiler.stop() |
| 40 | print(f"\n=== Performance trace: {func.__module__}.{func.__qualname__} ===") |
| 41 | print( |
| 42 | profiler.output_text( |
| 43 | color=color, |
| 44 | unicode=unicode, |
| 45 | show_all=show_all, |
| 46 | ) |
| 47 | ) |
| 48 | |
| 49 | return async_wrapper if is_coro else sync_wrapper |
| 50 | |
| 51 | return decorator |