in ppo_ewma/logger.py [0:0]
def profile(n):
"""
Usage:
@profile("my_func")
def my_func(): code
"""
def decorator_with_name(func, name):
@wraps(func)
def func_wrapper(*args, **kwargs):
with profile_kv(name):
return func(*args, **kwargs)
return func_wrapper
if callable(n):
return decorator_with_name(n, n.__name__)
elif isinstance(n, str):
return partial(decorator_with_name, name=n)
else:
raise NotImplementedError(
"profile should be called as either a bare decorator"
" or with a string (profiling name of a function) as an argument"
)