Python Class Based Decorator with parameters that can decorate a method or a function

Solution 1:

You don't need to mess around with descriptors. It's enough to create a wrapper function inside the __call__() method and return it. Standard Python functions can always act as either a method or a function, depending on context:

class MyDecorator(object):
    def __init__(self, argument):
        self.arg = argument

    def __call__(self, fn):
        @functools.wraps(fn)
        def decorated(*args, **kwargs):
            print "In my decorator before call, with arg %s" % self.arg
            result = fn(*args, **kwargs)
            print "In my decorator after call, with arg %s" % self.arg
            return result
        return decorated

A bit of explanation about what's going on when this decorator is used like this:

@MyDecorator("some other func!")
def some_other_function():
    print "in some other function!"

The first line creates an instance of MyDecorator and passes "some other func!" as an argument to __init__(). Let's call this instance my_decorator. Next, the undecorated function object -- let's call it bare_func -- is created and passed to the decorator instance, so my_decorator(bare_func) is executed. This will invoke MyDecorator.__call__(), which will create and return a wrapper function. Finally this wrapper function is assigned to the name some_other_function.

Solution 2:

You're missing a level.

Consider the code

class Foo(object):
    @MyDecorator("foo baby!")
    def bar(self):
        print "in bar!"

It is identical to this code

class Foo(object):
    def bar(self):
        print "in bar!"
    bar = MyDecorator("foo baby!")(bar)

So MyDecorator.__init__ gets called with "foo baby!" and then the MyDecorator object gets called with the function bar.

Perhaps you mean to implement something more like

import functools

def MyDecorator(argument):
    class _MyDecorator(object):
        def __init__(self, fn):
            self.fn = fn

        def __get__(self, obj, type=None):
            return functools.partial(self, obj)

        def __call__(self, *args, **kwargs):
            print "In my decorator before call, with arg %s" % argument
            self.fn(*args, **kwargs)
            print "In my decorator after call, with arg %s" % argument

    return _MyDecorator

Solution 3:

In your list of types of decorators, you missed decorators that may or may not take arguments. I think this example covers all your types except "function style decorators (wrapping a function)"

class MyDecorator(object):

    def __init__(self, argument):
        if hasattr('argument', '__call__'):
            self.fn = argument
            self.argument = 'default foo baby'
        else:
            self.argument = argument

    def __get__(self, obj, type=None):
        return functools.partial(self, obj)

    def __call__(self, *args, **kwargs):
        if not hasattr(self, 'fn'):
            self.fn = args[0]
            return self
        print "In my decorator before call, with arg %s" % self.argument
        self.fn(*args, **kwargs)
        print "In my decorator after call, with arg %s" % self.argument


class Foo(object):
    @MyDecorator("foo baby!")
    def bar(self):
        print "in bar!"

class Bar(object):
    @MyDecorator
    def bar(self):
        print "in bar!"

@MyDecorator
def add(a, b):
    print a + b