在 MRO 中装饰顶部功能

Posted 2023-02-23

【中文标题】在 MRO 中装饰顶部功能

【英文标题】：Decorating top function in MRO

【发布时间】：2021-05-03 03:23:43

【问题描述】：

如何修饰类继承中的最后一个函数？

如果我装饰一个超类函数，子类函数会覆盖装饰器。 我想知道是否有一种巧妙的方法可以自动装饰 MRO 中的顶部功能。

def wrapper(f):
    def _wrap(*args, **kwargs):
        print("In wrapper")
        return f(*args, **kwargs)
    return _wrap


class A:

    @wrapper
    def f(self):
        print("In class A")


class B(A):
    def f(self):
        print("In class B")


if __name__ == '__main__':
    a = A()
    b = B()
    print("Calling A:")
    a.f()
    print("Calling B:")
    b.f()

这是输出。正如预期的那样，B.f() 不会调用包装器，尽管我希望它这样做。

Calling A:
In wrapper
In class A

Calling B:
In class B

这是我迄今为止尝试过的。一个包含所有装饰器并在类实例化期间注入它们的元类。

from abc import ABCMeta


class WrapperMetaClass(ABCMeta):
    def __init__(cls, *args, **kwargs):
        wrappers_dict = getattr(cls, "_wrappers")
        for attr_name in dir(cls):
            if attr_name not in wrappers_dict:
                continue
            else:
                wrapper = wrappers_dict[attr_name]
            attr = getattr(cls, attr_name)
            if not hasattr(attr, '__call__'):
                raise Exception("What you're trying to wrap is not a function!")
            attr = wrapper(attr)
            setattr(cls, attr_name, attr)

        super().__init__(*args, **kwargs)

这行得通：

class A(metaclass=WrapperMetaClass):
    _wrappers = 
        "f": wrapper
    

    def f(self):
        print("In class A")


class B(A):
    def f(self):
        print("In class B")

输出是我想要的。

Calling A:
In wrapper
In class A

Calling B:
In wrapper
In class B

但是，这会遇到不同的问题。如果 B 不覆盖 f，则元类将 A.f() 包装两次。这是有道理的，因为 A 和 B 都继承了 WrapperMetaClass，所以先包装 A.f()，然后再包装 B.f()。

class A(metaclass=WrapperMetaClass):
    _wrappers = 
        "f": wrapper
    

    def f(self):
        print("In class A")


class B(A):
    pass

输出变成：

Calling A:
In wrapper
In class A

Calling B:
In wrapper
In wrapper
In class A

我不知道我还能做什么。

【问题讨论】：

如果B.f 调用A.f，你希望它表现得像A.f 被包裹，还是像A.f 未被包裹？

for attr_name in dir(cls): 如果要遍历类属性并修改它们，为什么不使用__new__()？

我只希望 B.f 被包裹起来。想想装饰器添加线程安全功能的用例，即锁。我希望 B.f 使用锁。 A.f 也使用锁是没有意义的。

@OlvinR​​oght 这实际上可行，我会尝试。

@IvanTsenov，这实际上是它应该如何工作的：D

【参考方案1】：

是的，我记得曾经遇到过一次或两次 - 你走在正确的轨道上。

但首先要做的事情是：如果您的“包装器”中的逻辑是 可以放在基类的方法中，然后分解方法 在较小的任务中，并且有一个“方法槽”系统比这个更可取， 正如user 2357112 supports monica 放入cmets。如果你发现你真的需要或者更喜欢装饰器，下面是完整的代码

class A:
    def do_things(self):
        create_connection()  # <- this is the code you'd are putting in the wrapper in the other approach
        do_thing_1()

class B(A):
    def do_things(self):
        # here we have to do thing_1 and thing_2, but
        # the connection is created in the superclass method... 
        # this could be the origin of your question
        
# Refactor to:

class A:
    def do_things(self):
        create_connection()
        self.internal_do_things()
        
    def internal_do_things(self):
        do_thing_1()
        
class B(A):
    def internal_do_things(self):
        super().internal_do_things()
        do_thing_2()

所以，经典继承和 OO 解决了这个问题

---

如果您需要装饰器：

要做的事情是让装饰器本身，“包装器”，得到 一种“知道”它是否已经在外部方法（即方法 在调用super()) 的子类中，并充当透明的 在这种情况下是包装器。

当我们想要一个强大的解决方案时，它会变得更加复杂： 一个可用于同一类中不同方法的包装器， 如果同时调用它们也不会感到困惑 （在不同的线程中，或者一个方法调用另一个方法， 不是super()，它应该触发包装器）。

最后，其机制足够复杂，以至于 他们不应该妨碍你的实际包装 - 所以， 理想情况下，它们应该自己构建为装饰器，这将 装饰你的包装。

[几小时后] 所以，很抱歉，如果它看起来不“整洁” - 结果是实现上述内容比我最初想象的要复杂一些 - 我们需要一个中间装饰器级别（在代码中称为 meta_wrapper_applier），以便元类每次重新声明方法时都可以重新包装方法。

希望代码和变量名中的cmets足以理解思路：

from abc import ABCMeta
from functools import wraps
import threading


class WrapperMetaClass(ABCMeta):
    def __init__(cls, name, bases, ns, **kw):
        super().__init__(name, bases, ns, **kw)

        # Get the wrapped methods for all the superclasses
        super_wrappers = 
        for supercls in cls.__mro__[::-1]:
            super_wrappers.update(supercls.__dict__.get("_wrappers", ))

        # unconditionally install a wrappers dict for each subclass:
        sub_wrappers =  cls._wrappers = 

        for attrname, attr in ns.items():
            if attrname in super_wrappers:
                # Applies the wrapper in the baseclass to the subclass method:
                setattr(cls, attrname, super_wrappers[attrname]._run_once_wrapper(attr))
            elif hasattr(attr, "_run_once_wrapper"):
                # Store the wrapper information in the cls for use of the subclasses:
                sub_wrappers[attrname] = attr


def run_once_method_decorator(original_wrapper):
    re_entering_stacks = 

    # This is the callable used to place a wrapper on the original
    # method and on each overriden method.
    # All methods with the same name in the subclasses  will share the same original wrapper and the
    # "re_entering_stacks" data structure.
    def meta_wrapper_applier(raw_method):
        wrapped_method_in_subclass = None
        @wraps(original_wrapper)
        def meta_wrapper(*args, **kw):
            nonlocal wrapped_method_in_subclass

            # uses a plain list to keep track of re-entering the same-named method
            # in each thread:
            re_entering_stack = re_entering_stacks.setdefault(threading.current_thread(), [])
            re_entering = bool(re_entering_stack)
            try:
                re_entering_stack.append(1)
                if re_entering:
                    result = raw_method(*args, **kw)
                else:
                    if wrapped_method_in_subclass is None:
                        # Applies the original decorator lazily, and caches the result:
                        wrapped_method_in_subclass = original_wrapper(raw_method)

                    result = wrapped_method_in_subclass(*args, **kw)
            finally:
                re_entering_stack.pop()
            return result
        # registry = original_wrapper.__dict__.setdefault("_run_once_registry", )


        meta_wrapper._run_once_wrapper = meta_wrapper_applier
        return meta_wrapper

    return meta_wrapper_applier

# From here on, example code only;

@run_once_method_decorator
def wrapper(f):
    @wraps(f)
    def _wrap(*args, **kwargs):
        print("Entering wrapper")
        result = f(*args, **kwargs)
        print("Leaving wrapper\n")
        return result
    return _wrap

@run_once_method_decorator
def other_wrapper(f):
    @wraps(f)
    def _wrap(*args, **kwargs):
        print("Entering other wrapper")
        result = f(*args, **kwargs)
        print("Leaving other wrapper\n")
        return result
    return _wrap


class A(metaclass=WrapperMetaClass):

    @wrapper
    def f(self):
        print("In class A")

    def g(self):
        print("g in A")

class B(A):
    def f(self):
        print("In class B")
        super().f()

    @other_wrapper
    def g(self):
        print("g in B")
        super().g()

class C(B):
    def g(self):
        print("g in C")
        super().g()


if __name__ == '__main__':
    a = A()
    b = B()
    print("Calling A:")
    a.f()
    a.g()
    print("Calling B:")
    b.f()
    b.g()
    print("Calling C:")
    C().g()

输出：

Calling A:
Entering wrapper
In class A
Leaving wrapper

g in A
Calling B:
Entering wrapper
In class B
In class A
Leaving wrapper

Entering other wrapper
g in B
g in A
Leaving other wrapper

Calling C:
Entering other wrapper
g in C
g in B
g in A
Leaving other wrapper