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

Adam Parkin picture Adam Parkin · Feb 23, 2012 · Viewed 20.1k times · Source

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

Sven Marnach picture Sven Marnach · Feb 23, 2012

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
            fn(*args, **kwargs)
            print "In my decorator after call, with arg %s" % self.arg
        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.