How to decorate a class method with a callable class?

Question

I want functions in a class to store their returned values in some data structure. For this purpose I want to use a decorator:

results = []
instances = []

class A:
    def __init__(self, data):
        self.data = data

    @decorator
    def f1(self, a, b):
        return self.data + a + b

    @decorator
    def f2(self, a):
        return self.data + a + 1

x = A(1)
x.f1(1, 2)
x.f2(3)

print(results)

The question is, how to implement this decorator.

The main idea is the following:

class Wrapper:

    def __init__(self, func):
        self.func = func

    def __call__(self, *args):
        res = self.func(*args)
        results.append(res)
        instances.append(args[0])

def decorator(func):
    return Wrapper(func)

But then I am getting an error message:

TypeError: f1() missing 1 required positional argument: 'b'

This question is similar to what other people asked (How to decorate a method inside a class?, Python functools.wraps equivalent for classes), but it is not clear where to put @functools.wraps or to call @functools.update_wrapped() in my case.

I'm a bit unclear on whether want class-specific or instance-specific storage. For example, if you have 2 instances of `A`, let's call them `x` and `y`, do you want these to write to the same list or not? — amdex, Jul 05 '19 at 14:24
My question is formulated the way that all instances write to the same list. Actually, I am more interested to register class instances as well. I updated my question. — Dmitry K., Jul 05 '19 at 14:29

score 3 · Answer 1 · answered Jul 05 '19 at 14:44

I think this is what you're asking for.

In [8]: from functools import wraps 
   ...: results = [] 
   ...: instances = [] 
   ...:  
   ...: class A: 
   ...:     def __init__(self, data): 
   ...:         self.data = data 
   ...:  
   ...:     def decorator(f): 
   ...:         @wraps(f) 
   ...:         def inner(*args, **kwargs): 
   ...:             retval = f(*args, **kwargs) 
   ...:             results.append(retval) 
   ...:             return retval  
   ...:         return inner 
   ...:         
   ...:  
   ...:     @decorator 
   ...:     def f1(self, a, b): 
   ...:         return self.data + a + b 
   ...:  
   ...:     @decorator 
   ...:     def f2(self, a): 
   ...:         return self.data + a + 1 
   ...:                                                                                         

In [9]: a = A(3)                                                                                

In [10]: a.f1(2,3)                                                                              
Out[10]: 8

In [11]: results                                                                                
Out[11]: [8]

Decorator does not have to be in the class, but it's handy when you need to access the class instance itself. You can take it out without any code change.

In [8]: from functools import wraps 
   ...: results = [] 
   ...: instances = [] 

   ...: def decorator(f): 
   ...:         @wraps(f) 
   ...:         def inner(*args, **kwargs): 
   ...:             retval = f(*args, **kwargs) 
   ...:             results.append(retval) 
   ...:             return retval  
   ...:         return inner 
   ...:  
   ...: class A: 
   ...:     def __init__(self, data): 
   ...:         self.data = data 
   ...:  
   ...:         
   ...:  
   ...:     @decorator 
   ...:     def f1(self, a, b): 
   ...:         return self.data + a + b 
   ...:  
   ...:     @decorator 
   ...:     def f2(self, a): 
   ...:         return self.data + a + 1 
   ...:                                                                                         

In [9]: a = A(3)                                                                                

In [10]: a.f1(2,3)                                                                              
Out[10]: 8

In [11]: results                                                                                
Out[11]: [8]

will work the same.

score 3 · Answer 2 · answered Jul 05 '19 at 14:55

If you want to have Wrapper in class outside, you can implement __get__ method (descriptor protocol):

from functools import partial

class Wrapper:
    def __init__(self, func):
        self.func = func

    def __get__(self, obj, type=None):
        bound_f = partial(self.__call__, obj)
        return bound_f

    def __call__(self, other_self, *args, **kwargs):
        res = self.func(other_self, *args, **kwargs)
        results.append(res)
        return res

def decorator(func):
    return Wrapper(func)

results = []

class A:
    def __init__(self, data):
        self.data = data

    @decorator
    def f1(self, a, b):
        return self.data + a + b

    @decorator
    def f2(self, a):
        return self.data + a + 1

x = A(1)
x.f1(1, 2)
x.f2(3)

print(results)

Prints:

[4, 5]

How to decorate a class method with a callable class?

2 Answers2