How to decorate a method inside a class?

Question

I am attempting to decorate a method inside a class but python is throwing an error. My class looks like this:

from pageutils import formatHeader

class myPage(object):
   def __init__(self):
      self.PageName = ''

   def createPage(self):
      pageHeader = self.createHeader()

   @formatHeader   #<----- decorator
   def createHeader(self):
       return "Page Header ",self.PageName

if __name__=="__main__":
   page = myPage()
   page.PageName = 'My Page'
   page.createPage()

---

pageutils.py:

def formatHeader(fn):
   def wrapped():
       return '<div class="page_header">'+fn()+'</div>'
   return wrapped

Python throws the following error

self.createHeader()
TypeError: wrapped() takes no arguments (1 given)

Where am I goofing?

Solution 1

Python automatically passes the class instance as reference. (The self argument which is seen in all instance methods).

You could do:

def formatHeader(fn):
    def wrapped(self=None):
        return '<div class="page_header">'+fn(self)+'</div>'
    return wrapped

Solution 2

You are omitting the self parameter which is present in the undecorated function (createHeader in your case).

def formatHeader(fn):
    from functools import wraps
    @wraps(fn)
    def wrapper(self):
        return '<div class="page_header">'+fn(self)+'</div>'
    return wrapper

If you are unsure about the signature of the function you want to decorate, you can make it rather general as follows:

def formatHeader(fn):
    from functools import wraps
    @wraps(fn)
    def wrapper(*args, **kw):
        return '<div class="page_header">'+fn(*args, **kw)+'</div>'
    return wrapper

Solution 3

You can also decorate the method at runtime, but not at define time. This could be useful in the case where you don't have access to or don't want to edit the source code, for example.

In [1]: class Toy():
   ...:     def __init__(self):
   ...:         return
   ...:     def shout(self, s):
   ...:         print(s)
   ...:

In [2]: def decor(fn):
   ...:     def wrapper(*args):
   ...:         print("I'm decorated")
   ...:         return fn(*args)
   ...:     return wrapper
   ...:


In [4]:

In [4]: a=Toy()

In [5]: a.shout('sa')
sa

In [6]: a.shout=decor(a.shout)

In [7]: a.shout('sa')
I'm decorated
sa

Comments

[exhuma]

oh... and in "createHeader", you are returning a tuple. Do this instead: return "Page Header " + self.PageName

[exhuma]

technically you could also do def wrapped(self), unless if you want to use the decorator also outside of a class. But then the decorated functions need to deal with self gracefully!

[jorgeh]

+1 for from functools import wraps. This has the added benefit that it exposes the docstring from the wrapped function. There's no reason not to use it!

[Charlie Parker]

so fn is the same as fn(self,...) or what is the correct way to think about where self if?

[Charlie Parker]

I get that self is None...? weird?

[exhuma]

@CharlieParker I don't quite understand you question.fn is a variable which holds a reference to the wrapped function. fn(self, ...) would be the call to that referenced function. Does that answer your question?

[ady]

from functools import wraps also exposes the name of the function: wrapped_function.__name__ (apart from its documentation) docs.python.org/2/library/functools.html#functools.wraps