Demo of how to write a decorator with optional arguments.

decorate_with_args.py

#!/usr/bin/env python

"""
Demo of how to write a decorator with optional arguments.
"""

import functools
import inspect

def mydec(*dec_args, **dec_kwargs):

    # Decorator invoked with arguments, e.g., @mydec() or @mydec(1, 2, 3)
    if not dec_args or not inspect.isfunction(dec_args[0]):
        def decorator(func):
            @functools.wraps(func)
            def wrapper(*args, **kwargs):
                print 'decorator args: ', dec_args, dec_kwargs
                print 'function args:  ', args, kwargs
                return func(*args, **kwargs)
            return wrapper
        return decorator

    # Decorator invoked with no arguments, e.g. @mydec
    else:
        func = dec_args[0]
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            print 'no decorator args'
            print 'function args:  ', args, kwargs
            return func(*args, **kwargs)
        return wrapper

# Same as above, but using the wrapt package:
import wrapt

def mydecw(*dec_args, **dec_kwargs):
    if not dec_args or not inspect.isfunction(dec_args[0]):
        @wrapt.decorator
        def wrapper(wrapped, instance, args, kwargs):
            print 'decorator args: ', dec_args, dec_kwargs
            print 'function args:  ', args, kwargs
            return wrapped(*args, **kwargs)
        return wrapper
    else:
        @wrapt.decorator
        def wrapper(wrapped, instance, args, kwargs):
            print 'no decorator args'
            print 'function args:  ', args, kwargs
            return wrapped(*args, **kwargs)
        return wrapper(dec_args[0])

if __name__ == '__main__':
    print 'implemented with functools:'

    @mydec(1, 2, 3)
    def foo(a):
        return a**2
    print foo(2)

    @mydec(1, 2, 3, x=4)
    def qux(a):
        return a**2
    print qux(2)

    @mydec()
    def bar(a):
        return a**2
    print bar(2)

    @mydec
    def baz(a):
        return a**2
    print baz(2)

    print 'implemented with wrapt:'

    @mydecw(1, 2, 3)
    def foo(a):
        return a**2
    print foo(2)

    @mydecw(1, 2, 3, x=4)
    def qux(a):
        return a**2
    print qux(2)

    @mydecw()
    def bar(a):
        return a**2
    print bar(2)

    @mydecw
    def baz(a):
        return a**2
    print baz(2)