Making decorators with optional arguments [duplicate]

The cleanest way I know of for doing this is the following:

import functools


def decorator(original_function=None, optional_argument1=None, optional_argument2=None, ...):

    def _decorate(function):

        @functools.wraps(function)
        def wrapped_function(*args, **kwargs):
            ...

        return wrapped_function

    if original_function:
        return _decorate(original_function)

    return _decorate

Explanation

When the decorator is called with no optional arguments like this:

@decorator
def function ...

The function is passed as the first argument and decorate returns the decorated function, as expected.

If the decorator is called with one or more optional arguments like this:

@decorator(optional_argument1='some value')
def function ....

Then decorator is called with the function argument with value None, so a function that decorates is returned, as expected.

Python 3

Note that the decorator signature above may be improved with Python 3-specific *, syntax to enforce safe use of keyword arguments. Simply replace the signature of the outermost function with:

def decorator(original_function=None, *, optional_argument1=None, optional_argument2=None, ...):

Through the help of the answers here and elsewhere and a bunch of trial and error I've found that there is actually a far easier and generic way to make decorators take optional arguments. It does check the args it was called with but there isn't any other way to do it.

The key is to decorate your decorator.

Generic decorator decorator code

Here is the decorator decorator (this code is generic and can be used by anyone who needs an optional arg decorator):

def optional_arg_decorator(fn):
    def wrapped_decorator(*args):
        if len(args) == 1 and callable(args[0]):
            return fn(args[0])

        else:
            def real_decorator(decoratee):
                return fn(decoratee, *args)

            return real_decorator

    return wrapped_decorator

Usage

Using it is as easy as:

1. Create a decorator like normal.
2. After the first target function argument, add your optional arguments.
3. Decorate the decorator with optional_arg_decorator

Example:

@optional_arg_decorator
def example_decorator_with_args(fn, optional_arg = 'Default Value'):
    ...
    return fn

Test cases

For your use case:

So for your case, to save an attribute on the function with the passed-in method name or the __name__ if None:

@optional_arg_decorator
def register_method(fn, method_name = None):
    fn.gw_method = method_name or fn.__name__
    return fn

Add decorated methods

Now you have a decorator that is usable with or without args:

@register_method('Custom Name')
def custom_name():
    pass

@register_method
def default_name():
    pass

assert custom_name.gw_method == 'Custom Name'
assert default_name.gw_method == 'default_name'

print 'Test passes :)'

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...')