Making decorators with optional arguments

Question

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 1

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, ...):