import cProfile
import io
import pstats
from functools import wraps
from typing import Callable, Optional
from elmo.logger import Logger as log # Or replace with `import logging`
[docs]
def function_profiler(
_func: Optional[Callable] = None, *, sortby: str = "cumtime", top_n: int = 50
):
"""
A decorator to profile function performance. Can be used with or without arguments.
Usage:
@function_profiler
def my_func(): ...
@function_profiler(sortby='tottime', top_n=25)
def another_func(): ...
:param _func: Internal placeholder for decorator usage without parameters.
:param sortby: str - Sorting key for profiling ('cumtime', 'tottime', etc).
:param top_n: int - Number of top functions to show in profile output.
"""
def decorator(func: Callable):
@wraps(func)
def wrapper(*args, **kwargs):
profiler = cProfile.Profile()
profiler.enable()
result = func(*args, **kwargs)
profiler.disable()
s = io.StringIO()
ps = pstats.Stats(profiler, stream=s).sort_stats(sortby)
ps.print_stats(top_n)
log.message(f"📊 Profiling results for {func.__name__}:\n{s.getvalue()}")
return result
return wrapper
# Handle both @decorator and @decorator(...)
if _func is None:
return decorator
else:
return decorator(_func)
[docs]
def function_profiler_old(sortby="cumtime", top_n=50):
"""
Decorator to profile a function and log its performance.
:param sortby: str - Sorting criteria for profiling results ('cumtime' or 'tottime').
:param top_n: int - Number of top entries to display in the profiling results.
"""
def decorator(func):
@wraps(func)
def wrapper(*args, **kwargs):
profiler = cProfile.Profile()
profiler.enable() # Start profiling
# Execute the wrapped function
result = func(*args, **kwargs)
profiler.disable() # Stop profiling
# Process profiling results
s = io.StringIO()
ps = pstats.Stats(profiler, stream=s).sort_stats(sortby)
ps.print_stats(top_n) # Print the top N entries
# Log the profiling results (replace with your logging mechanism)
log.message(f"Profiling results for {func.__name__}:\n{s.getvalue()}")
return result
return wrapper
return decorator