如何装饰__enter__和__exit__实例方法？解决装饰器未生效问题

Python的特殊方法（如__enter__、__exit__）是通过类层级的查找机制调用的，而非直接读取实例属性。当执行with语句时，解释器会去实例所属的类（这里是contMgr）中查找这两个方法，而不是使用你绑定到实例mgr上的装饰后方法，所以你的wrapper没有被触发。

下面提供几种可行的修复方式：

方案1：在类定义时直接装饰方法 #

直接在类内部给__enter__和__exit__添加装饰器，这是最直接的方式：

def deco(func):
    def wrapper(*args,**kwargs):
        print("In wrapper")
        result = func(*args, **kwargs)
        return result
    
    return wrapper

class contMgr():
    def __init__(self):
        print("In __init__")

    @deco  # 直接装饰__enter__
    def __enter__(self):
        print("In __enter__")
        return self

    @deco  # 直接装饰__exit__
    def __exit__(self, exc_type, exc_val, exc_tb):
        print("In __exit__")

    def __call__(self, *args, **kwargs):
        print("In __call__")

    def brol():
        print("brol")

@deco
def test():
    print("Blah, blah, blah")

mgr = contMgr()
print("------------------------")
with mgr:
    print("------------------------")
    test()
    print("------------------------")

运行后输出符合预期：

In __init__
------------------------
In wrapper
In __enter__
------------------------
In wrapper
Blah, blah, blah
------------------------
In wrapper
In __exit__

方案2：动态修改类的方法（而非实例） #

如果不想修改类定义，也可以在实例化后修改类的对应方法，这样所有实例都会生效：

def deco(func):
    def wrapper(*args,**kwargs):
        print("In wrapper")
        result = func(*args, **kwargs)
        return result
    
    return wrapper

class contMgr():
    def __init__(self):
        print("In __init__")

    def __enter__(self):
        print("In __enter__")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("In __exit__")

    def __call__(self, *args, **kwargs):
        print("In __call__")

    def brol():
        print("brol")

@deco
def test():
    print("Blah, blah, blah")

mgr = contMgr()
# 修改类的__enter__和__exit__方法，而非实例属性
contMgr.__enter__ = deco(contMgr.__enter__)
contMgr.__exit__ = deco(contMgr.__exit__)
print("------------------------")
with mgr:
    print("------------------------")
    test()
    print("------------------------")

此方案同样能得到预期输出，核心是修改类的属性而非单个实例。

方案3：使用类装饰器 #

如果需要批量给类的特殊方法加装饰，可使用类装饰器实现：

def deco(func):
    def wrapper(*args,**kwargs):
        print("In wrapper")
        result = func(*args, **kwargs)
        return result
    
    return wrapper

# 类装饰器：给目标类的__enter__和__exit__方法加装饰
def decorate_special_methods(cls):
    cls.__enter__ = deco(cls.__enter__)
    cls.__exit__ = deco(cls.__exit__)
    return cls

@decorate_special_methods
class contMgr():
    def __init__(self):
        print("In __init__")

    def __enter__(self):
        print("In __enter__")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("In __exit__")

    def __call__(self, *args, **kwargs):
        print("In __call__")

    def brol():
        print("brol")

@deco
def test():
    print("Blah, blah, blah")

mgr = contMgr()
print("------------------------")
with mgr:
    print("------------------------")
    test()
    print("------------------------")

这种方式更灵活，适合需要给多个类的特殊方法统一加装饰的场景。

内容的提问来源于stack exchange，提问作者Walther Hautermann