Saturday, December 20, 2008

How to decorate __init__ or another python class method

I want to decorate a class' __init__ method to magically attach my own instance members to the object being __init__'d and then do something else at the end of the __init__ call. So I need to be able to access to the original __init__'s 'self' reference.

I want code that looks like:

class Foo:
@my_decorator
def __init__(self,arg1, arg2):
self.arg1 = arg1
self.arg2 = arg2

To magically transform into something like:

class Foo:
def __init__(self,arg1, arg2):
#entrypoint
self.instance_member_inserted_by_decorator = 1

self.arg1 = arg1
self.arg2 = arg2

#exitpoint
print "I'm dancing on your mother's grave"


The tricky part of this is writing a decorator such that you can access the 'self' of the object being created. My first implementation looked like this:


class Dec:
def __init__(self,f):
self.f = f
def __call__(self):
print 'Entry point in decorator'
self.f(self)
print 'Exit point in decorator'

class Foo:
@Dec
def __init__(self):
print ' in originalFoo.__init__'
print ' self.__class__.__name__ ==',self.__class__.__name__
self.var1 = 'var1'

f = Foo()
print "f.__class__.__name__ ==",f.__class__.__name__
print "f.__init__.__class__.__name__ ==",f.__init__.__class__.__name__
print 'f.__dict__.has_key("var1") == ',f.__dict__.has_key("var1")


This prints:

Entry point in decorator
in originalFoo.__init__
self.__class__.__name__ == Dec
Exit point in decorator
f.__class__.__name__ == Foo
f.__init__.__class__.__name__ == Dec
f.__dict__.has_key("var1") == False


That implementation edits the Dec object instead of the Foo object. It also screws up the original __init__ call so that it doesn't initialize the Foo object at all!

In order to wrapper the __init__ call, you need to use decorators with arguments. This will give you access to the intended self object:

class Dec:
def __call__(self,f):
def wrap(init_self,*args,**kwargs):
print 'Entry point in decorator'
init_self.inserted_during_decoration = 1
f(init_self,*args,**kwargs)
print 'Exit point in decorator'
return wrap

class Foo:
@Dec()
def __init__(self,var1):
print ' in originalFoo.__init__'
print ' self.__class__.__name__ ==',self.__class__.__name__
self.var1 = var1

f = Foo('my_var1')
print "f.__class__.__name__ ==",f.__class__.__name__
print "f.__init__.__class__.__name__ ==",f.__init__.__class__.__name__
print 'f.__dict__.has_key("var1") == ',f.__dict__.has_key("var1")
print 'f.inserted_during_decoration ==',f.inserted_during_decoration


This prints out:'

Entry point in decorator
in originalFoo.__init__
self.__class__.__name__ == Foo
Exit point in decorator
f.__class__.__name__ == Foo
f.__init__.__class__.__name__ == instancemethod
f.__dict__.has_key("var1") == True
f.inserted_during_decoration == 1