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