orokusaki orokusaki -4 years ago 36
Python Question

Making decorators with optional arguments



from functools import wraps

def foo_register(method_name=None):
"""Does stuff."""
def decorator(method):
if method_name is None:
method.gw_method = method.__name__
else:
method.gw_method = method_name
@wraps(method)
def wrapper(*args, **kwargs):
method(*args, **kwargs)
return wrapper
return decorator


Example: The following decorates
my_function
with
foo_register
instead of ever making it to
decorator
.

@foo_register
def my_function():
print('hi...')


Example: The following works as expected.

@foo_register('say_hi')
def my_function():
print('hi...')


If I want it to work correctly in both applications (one using
method.__name__
and one passing the name in), I have to check inside of
foo_register
to see if the first argument is a decorator, and if so, I have to:
return decorator(method_name)
(instead of
return decorator
). This sort of "check to see if it's a callable" seems very hackish. Is there a nicer way to create a multi-use decorator like this?

P.S. I already know that I can require the decorator to be called, but that's not a "solution". I want the API to feel natural. My wife loves decorating, and I don't want to ruin that.

Answer Source

Glenn - I had to do it then. I guess I'm glad that there is not a "magic" way to do it. I hate those.

So, here's my own answer (method names different than above, but same concept):

from functools import wraps

def register_gw_method(method_or_name):
    """Cool!"""
    def decorator(method):
        if callable(method_or_name):
            method.gw_method = method.__name__
        else:
            method.gw_method = method_or_name
        @wraps(method)
        def wrapper(*args, **kwargs):
            method(*args, **kwargs)
        return wrapper
    if callable(method_or_name):
        return decorator(method_or_name)
    return decorator

Example usage (both versions work the same):

@register_gw_method
def my_function():
    print('hi...')

@register_gw_method('say_hi')
def my_function():
    print('hi...')
Recommended from our users: Dynamic Network Monitoring from WhatsUp Gold from IPSwitch. Free Download