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

Question

I've seen many examples of Python decorators that are:

• function style decorators (wrapping a function)
• class style decorators (implementing __init__, __get__, and __call__)
• decorators which do not take arguments
• decorators which take arguments
• decorators which are "method friendly" (ie can decorate a method in a class)
• decorators which are "function friendly" (can decorate a plain function
• decorators that can decorate both methods and functions

But I've never seen a single example which can do all of the above, and I'm having trouble synthesizing from various answers to specific questions (such as this one, this one, or this one (which has one of the best answers I've ever seen on SO)), how to combine all of the above.

What I want is a class-based decorator which can decorate either a method or a function, and that takes at least one additional parameter. Ie so that the following would work:

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

    def __get__(self, ....):
        # voodoo magic for handling distinction between method and function here

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


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


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

some_other_function()
Foo().bar()

And I would expect to see:

In my decorator before call, with arg some other func!
in some other function!
In my decorator after call, with arg some other func!
In my decorator before call, with arg foo baby!
in bar!
In my decorator after call, with arg foo baby!

Edit: if it matters, I'm using Python 2.7.

Answer 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.

Answer 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

Answer 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