How to create a Python decorator that can be used either with or without parameters?

前端 未结 13 1639
死守一世寂寞
死守一世寂寞 2020-11-30 21:26

I\'d like to create a Python decorator that can be used either with parameters:

@redirect_output(\"somewhere.log\")
def foo():
    ....

or

13条回答
  •  长情又很酷
    2020-11-30 21:29

    You need to detect both cases, for example using the type of the first argument, and accordingly return either the wrapper (when used without parameter) or a decorator (when used with arguments).

    from functools import wraps
    import inspect
    
    def redirect_output(fn_or_output):
        def decorator(fn):
            @wraps(fn)
            def wrapper(*args, **args):
                # Redirect output
                try:
                    return fn(*args, **args)
                finally:
                    # Restore output
            return wrapper
    
        if inspect.isfunction(fn_or_output):
            # Called with no parameter
            return decorator(fn_or_output)
        else:
            # Called with a parameter
            return decorator
    

    When using the @redirect_output("output.log") syntax, redirect_output is called with a single argument "output.log", and it must return a decorator accepting the function to be decorated as an argument. When used as @redirect_output, it is called directly with the function to be decorated as an argument.

    Or in other words: the @ syntax must be followed by an expression whose result is a function accepting a function to be decorated as its sole argument, and returning the decorated function. The expression itself can be a function call, which is the case with @redirect_output("output.log"). Convoluted, but true :-)

提交回复
热议问题