CODEWITHSUNDEEP
pythondecorator
gath
gath

Reputation: 25452

How to decorate a method inside a class?

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?

Upvotes: 74

Views: 53143

Answers (3)

exhuma
exhuma

Reputation: 21697

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

Upvotes: 47

Sibbs Gambling
Sibbs Gambling

Reputation: 20345

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

Upvotes: 2

krawyoti
krawyoti

Reputation: 20135

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

Upvotes: 68

Related Questions