Question

I have a class that looks like this:

class Test(object):
    data = {}

    @add_to_data(1)
    def method_1(self, x, y):
        pass

    @add_to_data(1, 2, 3)
    def method_2(self, x, y):
        pass

    @add_to_data(5, 6, 7)
    def method_3(self, x, y):
        pass

I want it so that at ‘import time’ the data atribute of the class is:

>>> Test.data
{'method_1': [1], 'method_2': [1, 2, 3], 'method_3': [5, 6, 7]}

Currently I have this (data is a defaultdict):

def add_to_data(*items):
    def decorator(func):
        for item in items:
            name = func.__name__
            cls = ??
            cls.data[name].append(item)
        return func
    return decorator

But I don’t know how to get the class object. Any help?

Answer 1

It isn’t particularly clean, but you can do this using sys._getframe() to inspect the calling scope (the decorator will be called in the context of the class definition, where data is a local variable.

So something like the following should work:

import sys

def add_to_data(*items):
    def decorator(func):
        name = func.__name__
        data = sys._getframe(1).f_locals.get('data')
        if data is not None:
            data.setdefault(name, []).extend(items)
        return func
    return decorator

If you don’t want to depend on sys._getframe or want something more explicit, you could make your @add_to_data function decorator simply mark the methods with an appropriate function attribute and then use a class decorator to build the data dictionary from all the marked methods.