Is there a way apply a decorator to a Python method that needs informations about the class?

Question

When you decorate a method, it is not bound yet to the class, and therefor doesn't have the im_class attribute yet. I looking for a way to get the information about the class inside the decorator. I tried this:

import types

def decorator(method):

    def set_signal(self, name, value):
        print name
        if name == 'im_class':
            print "I got the class"

    method.__setattr__ = types.MethodType(set_signal, method)

    return method


class Test(object):
    @decorator
    def bar(self, foo):
        print foo

But it doesn't print anything.

I can imagine doing this:

class Test(object):
    @decorator(klass=Test)
    def bar(self, foo):
        print foo

But if I can avoid it, it would make my day.

Answer 1

__setattr__ is only called on explicit object.attribute = assignments; building a class does not use attribute assignment but builds a dictionary (Test.__dict__) instead.

To access the class you have a few different options though:

1. Use a class decorator instead; it'll be passed the completed class after building it, you could decorate individual methods on that class by replacing them (decorated) in the class. You could use a combination of a function decorator and a class decorator to mark which methods are to be decorated:

   def methoddecoratormarker(func):
       func._decorate_me = True
       return func
   
   def realmethoddecorator(func):
       # do something with func. 
       # Note: it is still an unbound function here, not a method!
       return func
   
   def classdecorator(klass):
       for name, item in klass.__dict__.iteritems():
           if getattr(item, '_decorate_me', False):
               klass.__dict__[name] = realmethoddecorator(item)
   

   You could use a metaclass instead of a class decorator to achieve the same, of course.

2. Cheat, and use sys._getframe() to retrieve the class from the calling frame:

   import sys
   
   def methoddecorator(func):
        callingframe = sys._getframe(1)
        classname = callingframe.f_code.co_name
   

   Note that all you can retrieve is the name of the class; the class itself is still being built at this time. You can add items to callingframe.f_locals (a mapping) and they'll be made part of the new class object.

3. Access self whenever the method is called. self is a reference to the instance after all, and self.__class__ is going to be, at the very least, a sub-class of the original class the function was defined in.

Answer 2

My strict answer would be: It's not possible, because the class does not yet exist when the decorator is executed.

The longer answer would depend on your very exact requirements. As I wrote, you cannot access the class if it does not yet exists. One solution would be, to mark the decorated method to be "transformed" later. Then use a metaclass or class decorator to apply your modifications after the class has been created.

Another option involves some magic. Look for the implementation of the implements method in zope.interfaces. It has some access to the information about the class which is just been parsed. Don't know if it will be enough for your use case.

Answer 3

You might want to take a look at descriptors. They let you implement a __get__ that is used when an attribute is accessed, and can return different things depending on the object and its type.

Answer 4

Use method decorators to add some marker attributes to the interesting methods, and use a metaclass which iterates over the methods, finds the marker attributes, and does the logic. The metaclass code is run when the class is created, so it has a reference to the newly created class.

class MyMeta(object):
  def __new__(...):
    ...
    cls = ...
    ... iterate over dir(cls), find methods having .is_decorated, act on them
    return cls
def decorator(f):
  f.is_decorated = True
  return f
class MyBase(object):
  __metaclass__ = MyMeta
class MyClass(MyBase):
  @decorator
  def bar(self, foo):
    print foo

If you worry about that the programmer of MyClass forgets to use MyBase, you can forcibly set the metaclass in decorator, by exampining the globals dicitionary of the caller stack frame (sys._getframe()).