clean approach to decorator that takes optional keyword arguments [closed]

Question

I write decorators a lot in my work. Sometimes you need to write "decorator that takes arguments", it usually turns out to be something like this:

def decorator(**kwargs):

    def _decorator(func):

        @wraps(func)
        def decorated(*func_args, **func_kwargs):
            # do something with kwargs here
            return func(*func_args, **func_kwargs)

        return decorated

    return _decorator

but in this case you'll always have to use your decorator like so:

@decorator()  # if you forget () here you'll get an exception
def myfunc():
    pass

So I come up with this snippet, IMO this is nicer to read/write and understand:

from functools import wraps, partial

def decorator(func=None, **kwargs):
    if func:

        @wraps(func)
        def decorated(*func_args, **func_kwargs):

            if kwargs.get("sayhi"):
                print("hi")

            return func(*func_args, **func_kwargs)

        return decorated

    else:
        return partial(decorator, **kwargs)

@decorator
def hello(name):
    print(f"hello {name}")


@decorator(sayhi=True)
def bye(name):
    print(f"bye {name}")

hello("rabbit")
# "hello rabbit"

bye("rabbit")
# hi
# bye rabbit

Answer 1

Nice little snippet! A couple of minor points:

• It's usually better to put the shorter block first in an if else. This stops the else statement being far from the if, making it clearer how the logic flows.

• When comparing to None, check if func is None or similar, rather than comparing to True/False.

• Although this is just a toy example, docstrings should always be used to show what the code does and returns.

• PEP8 spacing between imports, functions and code.

• Argument names could be better but without knowing the application it's hard to say what they should be.

from functools import partial, wraps


def decorator(func=None, **kwargs):
    """
    Decorator which optionally takes an argument sayhi (True or False).
    prints 'hi' if True, function is then returned as normal
    """
    if func is None:
        return partial(decorator, **kwargs)
    else:
        @wraps(func)
        def decorated(*func_args, **func_kwargs):
            if kwargs.get('sayhi'):
                print('hi')
            return func(*func_args, **func_kwargs)
        return decorated