def profile()

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"
        )