optional_arg_decorator.py

"""
Wraps a decorator such that it can be used directly instead of called if no parameters are used.
For example, a decorator wrapped like:

    @optional_arg_decorator
    def my_decorator(fn, option=None):
        ...

Could be used like:

    @my_decorator
    def some_fn(): ...

or:

    @my_decorator(option=123)
    def some_fn(): ...

"""

def optional_arg_decorator(fn):
    # Shamelessly stolen from StackOverflow
    # https://stackoverflow.com/questions/3888158/making-decorators-with-optional-arguments#comment65959042_24617244

    def wrapped_decorator(*args, **kwargs):
        if len(args) == 1 and callable(args[0]):
            return fn(*args)
        else:
            def real_decorator(decoratee):
                return fn(decoratee, *args, **kwargs)

            return real_decorator

    return wrapped_decorator