装饰器与上下文管理

学习目标

• 理解函数是对象以及装饰器的闭包结构。
• 用 functools.wraps 保留被装饰函数元数据。
• 用 with 和 contextlib 管理资源生命周期。

核心知识点

装饰器接收一个可调用对象并返回新的可调用对象，常用于日志、计时、权限等横切逻辑。上下文管理器实现 __enter__/__exit__，保证资源在离开 with 块时清理。

常用用法

最小装饰器

def announce(function):
    def wrapper(*args, **kwargs):
        result = function(*args, **kwargs)
        return f"result={result}"

    return wrapper


@announce
def add(left: int, right: int) -> int:
    return left + right


print(add(2, 3))
# 输出：result=5

装饰器语法等价于 add = announce(add)，定义顺序很重要。

functools.wraps

from functools import wraps


def trace(function):
    @wraps(function)
    def wrapper(*args, **kwargs):
        return function(*args, **kwargs)

    return wrapper


@trace
def greet(name: str) -> str:
    """返回问候语。"""
    return f"Hi, {name}"


print(greet.__name__, greet.__doc__)
# 输出：greet 返回问候语。

wraps 复制 __name__、__doc__ 等元数据，便于调试、文档和类型工具识别原函数。

带参数的装饰器

from functools import wraps


def repeat(times: int):
    def decorate(function):
        @wraps(function)
        def wrapper(*args, **kwargs):
            return [function(*args, **kwargs) for _ in range(times)]

        return wrapper

    return decorate


@repeat(2)
def identify(value: str) -> str:
    return value


print(identify("ok"))
# 输出：['ok', 'ok']

三层结构分别负责配置、接收函数、接收调用参数；装饰器配置应在定义时校验。

类作为装饰器

class CallCounter:
    def __init__(self, function):
        self.function = function
        self.count = 0

    def __call__(self, *args, **kwargs):
        self.count += 1
        return self.function(*args, **kwargs)


@CallCounter
def double(value: int) -> int:
    return value * 2


double(2)
print(double(3), double.count)
# 输出：6 2

类装饰器适合需要保存状态的场景，但要考虑如何暴露被装饰函数的元数据。

with 使用上下文管理器

from io import StringIO


buffer = StringIO()
with buffer as stream:
    stream.write("hello")
    content = stream.getvalue()
print(buffer.closed, content)
# 输出：True hello

with 把获取和释放资源绑定起来；文件、锁和临时目录都应优先使用此形式。

contextlib.contextmanager

from contextlib import contextmanager


@contextmanager
def labeled(label: str):
    print(f"enter:{label}")
    # 输出：enter:demo
    try:
        yield label.upper()
    finally:
        print(f"exit:{label}")
        # 输出：exit:demo


with labeled("demo") as value:
    result = value
print(result)
# 输出：DEMO

yield 前是进入逻辑，之后是退出逻辑；finally 保证异常时也清理。

不常用但需要知道

functools.lru_cache

from functools import lru_cache


@lru_cache(maxsize=2)
def square(value: int) -> int:
    return value * value


print(square(4), square.cache_info().hits)
# 输出：16 0

缓存要求参数可哈希，且函数应是无副作用的；数据变化后要考虑 cache_clear()。

contextlib.suppress

from contextlib import suppress


values = {"name": "Ada"}
with suppress(KeyError):
    del values["missing"]
print(values)
# 输出：{'name': 'Ada'}

只对明确可忽略的异常使用，避免把错误隐藏在宽泛的 suppress 中。

contextlib.ExitStack

from contextlib import ExitStack
from io import StringIO


with ExitStack() as stack:
    streams = [stack.enter_context(StringIO()) for _ in range(2)]
    streams[0].write("a")
print(len(streams), streams[0].closed)
# 输出：2 True

资源数量动态变化时，ExitStack 能统一管理退出顺序。

易混点

• 装饰器在定义时执行，包装函数在调用时执行。
• 不用 wraps 会丢失函数名和文档，某些框架也会因此无法识别签名。
• with 不会自动捕获所有异常，只负责进入和退出协议。
• 装饰器顺序从下到上应用，调用时从上到下包裹。
• 缓存只适合可重复、参数可哈希且结果稳定的函数。

课后小问

1. @trace 语法的等价写法是什么？ 答案：函数定义后执行 function = trace(function)。 解析：装饰器替换的是绑定到函数名的对象，调用者之后看到的是包装函数。

2. 为什么上下文管理器要把清理写在 finally？ 答案：这样正常离开和异常离开都会执行清理。 解析：资源释放不能依赖业务主体是否成功，否则异常路径会泄漏资源。

本节小结

装饰器封装调用前后行为，上下文管理器封装资源进出；二者都应保持边界小、异常语义明确、元数据可追踪。

快速回顾

• 函数包装：闭包、@decorator
• 元数据：functools.wraps
• 参数装饰器：三层函数结构
• 资源管理：with、__enter__、__exit__
• 工具：contextmanager、ExitStack、lru_cache