CODEWITHSUNDEEP
pythonmetaprogramming
priestc
priestc

Reputation: 35220

Modifying a class attribute via method decorator

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?

Upvotes: 3

Views: 1050

Answers (2)

James Henstridge
James Henstridge

Reputation: 43929

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.

Upvotes: 4

abarnert
abarnert

Reputation: 365895

You can't do that, because the class object doesn't exist yet.

You probably could do something similar by either creating a base class that has a @classmethod to use as a decorator in subclasses, or a decorator that stores things somewhere else together with a class decorator or metaclass that copies the storage into the class's data attribute later.

Well, really, without knowing what your actual goal is—and whether you're on Py2 or Py3 (and whether you only care about CPython or any Python)—it's hard to know what to suggest.

Upvotes: 1

Related Questions