decorator soup

gistfile1.py

import functools

def _decorate_method(f, name):
    @functools.wraps(f)
    def __wrapper(self, *args, **kwds):
        print 'wrapped a method', f, name
        return f(self, *args, **kwds)
    return __wrapper


def _decorate_class(cls, name):
    original_init = cls.__init__

    def __init_wrapper(self, *args, **kwds):
        print 'wrapped a class __init__', cls, name
        return original_init(self, *args, **kwds)
    cls.__init__ = __init_wrapper
    return cls

def _make_decorator(name, f):
    if isinstance(f, type):
        print 'decorating class', f, name
        return _decorate_class(f, name)
    else:
        print 'decorating method', f, name
        return _decorate_method(f, name)

def my_decorator(arg='some-default'):
    """
    decorator generator logic engine wrapper

    determines if the thing to be wrapped is a method or a class, then
    calls the appropriate decorator generator
    """
    if isinstance(arg, str):
        # requires partial application, because the return must be a
        # function that takes a function as its argument
        return functools.partial(_make_decorator, arg)
    else:
        return _make_decorator('some-default', arg)



## examples
@my_decorator
@my_decorator()
@my_decorator('not-bar')
class Foo(object):
    @my_decorator
    def something_a(self):
        pass
    
    @my_decorator()
    def something_b(self):
        pass
    
    @my_decorator('not-foo')
    def something_c(self):
        pass

print 'instantiation here'
f = Foo()
print 'calling methods'
f.something_a()
f.something_b()
f.something_c()

## output
# decorating method <function something_a at 0x10f4a1320> some-default
# decorating method <function something_b at 0x10f4a1410> some-default
# decorating method <function something_c at 0x10f4a1500> not-foo
# decorating class <class '__main__.Foo'> not-bar
# decorating class <class '__main__.Foo'> some-default
# decorating class <class '__main__.Foo'> some-default
# instantiation here
# wrapped a class __init__ <class '__main__.Foo'> some-default
# wrapped a class __init__ <class '__main__.Foo'> some-default
# wrapped a class __init__ <class '__main__.Foo'> not-bar
# calling methods
# wrapped a method <function something_a at 0x10f4a1320> some-default
# wrapped a method <function something_b at 0x10f4a1410> some-default
# wrapped a method <function something_c at 0x10f4a1500> not-foo