通过装饰器向类动态添加函数

Question

I'm trying to find a way to dynamically add methods to a class through decorator. The decorator i have look like:

def deco(target):

    def decorator(function):
        @wraps(function)
        def wrapper(self, *args, **kwargs):
            return function(*args, id=self.id, **kwargs)

        setattr(target, function.__name__, wrapper)
        return function

    return decorator

class A:
    pass

# in another module
@deco(A)
def compute(id: str):
    return do_compute(id)

# in another module
@deco(A)
def compute2(id: str):
    return do_compute2(id)

# **in another module**
a = A()
a.compute() # this should work
a.compute2() # this should work

My hope is the decorator should add the compute() function to class A, any object of A should have the compute() method. However, in my test, this only works if i explicitly import compute into where an object of A is created. I think i'm missing something obvious, but don't know how to fix it. appreciate any help!

Answer 1

I think this will be quite simpler using a decorator implemented as a class:

class deco:
    def __init__(self, cls):
        self.cls = cls

    def __call__(self, f):
        setattr(self.cls, f.__name__, f)
        return self.cls

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

@deco(A)
def compute(a_instance):
    print(a_instance.val)


A(1).compute()
A(2).compute()

outputs

1
2

But just because you can do it does not mean you should. This can become a debugging nightmare, and will probably give a hard time to any static code analyser or linter (PyCharm for example "complains" with Unresolved attribute reference 'compute' for class 'A' )

---

Why doesn't it work out of the box when we split it to different modules (more specifically, when compute is defined in another module )?

Assume the following:

a.py

print('importing deco and A')

class deco:
    def __init__(self, cls):
        self.cls = cls

    def __call__(self, f):
        setattr(self.cls, f.__name__, f)
        return self.cls

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

b.py

print('defining compute')

from a import A, deco

@deco(A)
def compute(a_instance):
    print(a_instance.val)

main.py

from a import A

print('running main')

A(1).compute()
A(2).compute()

If we execute main.py we get the following:

importing deco and A
running main
Traceback (most recent call last):
    A(1).compute()
AttributeError: 'A' object has no attribute 'compute'

Something is missing. defining compute is not outputted. Even worse, compute is never defined, let alone getting bound to A .

Why? because nothing triggered the execution of b.py . Just because it sits there does not mean it gets executed.

We can force its execution by importing it. Feels kind of abusive to me, but it works because importing a file has a side-effect: it executes every piece of code that is not guarded by if __name__ == '__main__ , much like importing a module executes its __init__.py file.

main.py

from a import A
import b

print('running main')

A(1).compute()
A(2).compute()

outputs

importing deco and A
defining compute
running main
1
2