装饰器与 functools

周五收盘后，你在私募自营盘的代码库里逛了一圈，发现满屏都是 @ 开头的行——@dataclass、@contextmanager、@lru_cache、@property，还有内部框架塞进来的 @trace、@timed。同事让你帮忙给一个慢得离谱的债券估值函数 bond_price 提速：估算 300ETF 篮子里上千只债时，同一组 (face, rate, n_periods) 被反复调用，profiler 报告说光这一个函数吃掉了 40% 的 CPU 时间。你知道答案应该是 @lru_cache，但「@ 到底是什么」这件事一直没正经拆开看过。这一课把它拆到底：装饰器（decorator）是「拿一个函数返回另一个函数」的函数，functools 则是给这类用法准备的工具箱——wraps、lru_cache、partial 都从同一个零件衍生出来。

@ 只是语法糖

先把魔法去掉。@deco 写在 def f(...): 上一行时，Python 做的事情和你手写 f = deco(f) ​一字不差​​。换句话说，deco 是一个普通函数，只不过它接受的参数和返回的值都是函数本身。

按这条退糖（desugar）路径从零写一个 trace：给任意函数加上「调用前不打印、调用后打印参数与返回值」的副作用。内层 wrapper 用 *args, **kwargs 收下任意签名，转发给原函数，再围绕这一次调用打一行日志：

def trace(fn):
    def wrapper(*args, **kwargs):
        result = fn(*args, **kwargs)
        print(f"调用 {fn.__name__}{args} -> {result}")
        return result
    return wrapper


# 写法 A：@ 语法糖
@trace
def add(a, b):
    return a + b


# 写法 B：手动赋值——与 A 完全等价
def add(a, b):
    return a + b
add = trace(add)


add(1, 2)   # 调用 add(1, 2) -> 3

两段写法的字节码是同一份。理解这件事之后，再看 @functools.lru_cache(maxsize=128)、@dataclass、@property，都不过是「同一根管子里跑着不同的 deco」。

functools.wraps：别擦掉函数身份

把上面的 trace 装上去之后，运行 print(add.__name__)——你会发现它变成了 'wrapper'，而不是 'add'。__doc__ 也被一并擦掉。调试器、tracing 工具、文档生成器都靠 __name__ 和 __doc__ 来定位函数；被装饰器一覆盖，整套日志里所有出错栈都长一个样：wrapper。

修这件事的标准做法是给内层 wrapper 再套一个装饰器——functools.wraps，它把原函数的 __name__、__doc__、__module__、__wrapped__ 等元信息搬到 wrapper 上：

from functools import wraps


def trace(fn):
    @wraps(fn)
    def wrapper(*args, **kwargs):
        result = fn(*args, **kwargs)
        print(f"调用 {fn.__name__}{args} -> {result}")
        return result
    return wrapper


@trace
def add(a, b):
    return a + b


add.__name__   # 'add'  (没加 @wraps(fn) 时是 'wrapper')

规则只有一句：​**​你写的每一个装饰器，内层 wrapper 上都要带 @wraps(fn)**​。这是可调试装饰器与不可调试装饰器之间唯一的差别。

functools.lru_cache：记忆化

lru_cache 解决的是另一类问题：把同一组参数算过的结果缓存（cache）起来，下次直接返回。前置条件两条——函数必须确定（deterministic），同一组参数永远算出同一个结果；参数必须可哈希（hashable），因为缓存键就是 (args, kwargs)。最经典的演示是递归 Fibonacci：没有缓存时是指数复杂度，加上一行装饰器就变成线性：

from functools import lru_cache


@lru_cache(maxsize=128)
def fib(n):
    return n if n < 2 else fib(n-1) + fib(n-2)


fib(30)   # 832040

不加缓存时，fib(30) 要重复展开约 2.7 亿次自递归，本机要跑好几秒；加上 @lru_cache(maxsize=128) 后立即返回 832040，因为每个 fib(k) 只算一次。maxsize 控制最多保留多少组结果，超出时按最近最少使用（least-recently-used, LRU）淘汰；functools.cache 是 maxsize=None 的无界版本，一句话提一下即可。

回到开头的 bond_price：只要它对同一组 (face, rate, n_periods) 永远算出同一个价，套上 @lru_cache(maxsize=128) 就一行解决，profiler 报告里那 40% 当场归零。

@lru_cache(maxsize=128) 这种写法叫​​带参数的装饰器​​（decorator factory）——lru_cache(maxsize=128) 先返回一个真正的装饰器，再贴到 fib 上。自己写一个带参数的装饰器需要再多一层嵌套，模板留到 3.3.2（Patterns & Tooling）再讲；这一课你只需要会用现成的。

functools.partial：偏函数

partial 不是装饰器，但和装饰器同属「函数操纵函数」这一族——把多参函数预先绑住几个参数，得到一个少参的新可调用对象。3.1.1 课里出现过的 P&L 函数 pnl 是个典型对象，回测时手续费（fee）经常被固定为 0：

from functools import partial


def pnl(p0, p1, shares, fee):
    return (p1 - p0) * shares - fee


pnl_no_fee = partial(pnl, fee=0.0)
pnl_no_fee(100.0, 101.5, 100)   # 150.0

等价的 lambda 写法是 lambda p0, p1, shares: pnl(p0, p1, shares, fee=0.0)——同样能用，多一行噪音。判别原则一句话：​​只固定参数​​用 partial，​​还要重排或加工参数​​就写 lambda。后者更灵活，代价是别人读你代码时要把这一行再读一遍。

你早就在用装饰器

回头看你这周的代码：@dataclass（第 1 课）、@contextlib.contextmanager（第 2 课），以及 Python 自带的 @property，本质上都是「拿函数（或类）返回另一个函数（或类）」的装饰器。区别只在它们藏的复杂度——dataclass 替你生成 __init__、__repr__；contextmanager 替你把一个 yield 翻译成 __enter__ / __exit__。这一课没让你写新东西，只是把你早就在用的零件拆开看了一次。

练习

def slow_square(n):
    time.sleep(0.01)
    return n * n

衔接下一课

到这里，你已经能给任意函数加横切关注点（cross-cutting concerns）：缓存、计时、tracing——都是一行装饰器。但 print 不是生产代码留痕的工具：它进不了日志聚合，关不掉，也分不出严重级别。下一课把 print 换成正经的 logging 模块，顺带把脚本的命令行入口、虚拟环境（venv）、以及 ruff / black / mypy 这套 CI 工具链一并接上——同一个 bond_price 模块，下周一就要作为命令行工具交付给风控组了。