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.

Malines

A "decorator that takes parameters" is just a function that takes the parameters and returns a decorator. – Trictrac

And why do you need to deal with methods and functions separately? Just pass all the arguments through. – Trictrac

@katrielalex, A method begins its life as a normal function and is stored on the class as one. When you look up a method it becomes a bound method, where the first argument to the function will be the instance on which you looked up the method. When you have objects that are instances of your own class rather than objects that are functions, they don't do this automatically. – Fernandafernande

@Trictrac there may be some very specified cases where you have to treat decoration for methods and "regular" functions differently. – Rata

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.

Selfhood

The key aspect to this I think OP was missing is that there is another level of callables: MyDecorator is called and its result is called, returning the object we store in the class/module (which is later called however many times). – Fernandafernande

@MikeGraham: I've added a fair bit of explanation to my post. – Selfhood

@MikeGraham: that's exactly right, I didn't understand the 2nd level of indirection. – Malines

Inside MyDecorator.__call__, functools.wraps is used to take care of some housekeeping details - copying the function name, docstring, and arguments list from the wrapped function (bare_func) to the wrapper function. There's a nice explanation of functools.wraps in the SO question "What does functools.wraps do?" – Ulphia

might want to return fn(*args, **kwargs) – Xanthochroid

@Xanthochroid Thanks, I've updated the answer. (Note that that part was copied from the question.) – Selfhood

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

Fernandafernande

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

Effeminacy

Only answer that doesn't use local functions or classes, which is exactly what I needed! – Lingam

Thanks Vinayak! Decorators with optional arguments that what I was looking for. – Livvi

This is truly brilliant! I've tested it and it does exactly what's needed. But can I ask you to kindly explain how it works. It is rather mysterious. At what point is get called and by whom? And why is the type argument needed and what is it? I understand partial from the functools doc, but get doc is very generic and I can't get my head around who is calling it during decoration and use of a decorated function. – Holtz

@BerndWechner: get will make MyDecorator a descriptor. Please refer: docs.python.org/3/howto/descriptor.html#descriptor-protocol f = Foo() f.bar() # This will invoke MyDecorator_object.__get__(f, Foo) – Effeminacy