3
\$\begingroup\$

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
\$\endgroup\$
  • 2
    \$\begingroup\$ Code Review requires real code. Do you have a concrete use case that is not a toy example, so that we can give you proper advice? \$\endgroup\$ – 200_success Oct 9 at 1:34
  • \$\begingroup\$ I agree with @200_success. This is borderline-hypothetical code; have a read through codereview.meta.stackexchange.com/questions/1981/… \$\endgroup\$ – Reinderien Oct 9 at 1:35
  • \$\begingroup\$ hmmm it isn't really a question i just want to share this with people, where should i post it? \$\endgroup\$ – rabbit.aaron Oct 9 at 3:42
  • \$\begingroup\$ @rabbit.aaron Github. \$\endgroup\$ – Mast Oct 9 at 9:38
  • \$\begingroup\$ I disagree, this is working code. though the real logic isn't displayed inside the code, it serves as a template for myself to write decorators. The purpose of the code is clearly stated in my post. I do think this question complies with most of the things mentioned in the guideline here. codereview.meta.stackexchange.com/questions/1954/… . if the code isn't working i would have posted it on stackoverflow. \$\endgroup\$ – rabbit.aaron Oct 9 at 14:00
3
\$\begingroup\$

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
\$\endgroup\$

Not the answer you're looking for? Browse other questions tagged or ask your own question.