进阶篇 / 第 10 章

第 10 章 · 迭代器、生成器与装饰器

迭代协议：for 循环背后的故事

你写了无数次 for 循环，但它究竟是怎么工作的？答案藏在一个叫「迭代器协议（iterator protocol）」的约定里。

可迭代对象 vs 迭代器

这条协议规定：

1. 可迭代对象调用 iter(obj)（内部即 obj.__iter__()）返回一个迭代器。
2. 迭代器调用 next(it)（内部即 it.__next__()）返回下一个元素。
3. 元素取完时，迭代器抛出 StopIteration 异常。

protocol_demo.py · 手动模拟 for 循环

data = ["a", "b", "c"]

# 普通写法
for x in data:
    print(x)

# 等价的展开写法
it = iter(data)        # 调 data.__iter__()，得到列表迭代器
while True:
    try:
        x = next(it)    # 调 it.__next__()
        print(x)
    except StopIteration:
        break           # 元素耗尽，for 内部就是靠这个结束

# 更简洁的官方等价写法（for 循环的内部逻辑）
it = iter(data)
for x in it:            # for 会自动捕获 StopIteration
    print(x)

哪些对象是可迭代的

下面这些对象都实现了迭代协议，所以能放进 for：

nums = [1, 2, 3]

print(hasattr(nums, "__iter__"), hasattr(nums, "__next__"))   # True False
it = iter(nums)
print(hasattr(it, "__iter__"), hasattr(it, "__next__"))       # True True

# 迭代器耗尽后再次遍历 = 空
print(list(it))   # []
print(list(it))   # []

自己写一个迭代器

range_step.py · 自定义迭代器类

class Countdown:
    """一个倒计时迭代器：从 start 每次递减 step，降到 stop 以下停止。"""

    def __init__(self, start: int, stop: int = 0, step: int = 1) -> None:
        if step == 0:
            raise ValueError("step 不能为 0")
        self.current = start
        self.stop = stop
        self.step = step

    def __iter__(self):
        # 迭代器的 __iter__ 必须返回自身
        return self

    def __next__(self) -> int:
        if (self.step > 0 and self.current <= self.stop) or \
           (self.step < 0 and self.current >= self.stop):
            raise StopIteration        # 必须抛这个异常才算结束
        value = self.current
        self.current -= self.step
        return value


# 使用
for n in Countdown(5):
    print(n, end=" ")          # 5 4 3 2 1
print()

for n in Countdown(10, 0, -3):
    print(n, end=" ")          # 10 7 4 1
print()

# 迭代器也是可迭代对象，可以用 list() 收集
print(list(Countdown(3)))     # [3, 2, 1]

# 反向来看：list 的迭代器可以转回 list
it = iter([1, 2, 3])
print(list(it), list(it))     # [1, 2, 3] []

可迭代对象的三种来源

在实践中，数据进入 for 循环有三条路：

1. 来自容器：list、dict、set——数据已经在内存里了，迭代只是走一遍。
2. 来自计算：range、x ** 2 for x in range(10)——数据是算出来的，不需要全部物化。
3. 来自外部：文件、socket、网络响应——数据是流式到达的，边到边处理。

后两种场景就是「惰性（lazy）」与「流式（streaming）」的用武之地，也就是下一节的主题。

生成器：把循环变成函数

生成器（generator）是 Python 里最优雅的语言机制之一。它解决的问题是：当你只需要逐个处理海量数据时，为什么必须先把所有数据放进内存？

普通函数 vs 生成器函数

两者的区别只有一处：函数体里出现了 yield，整个函数就自动变成生成器函数（generator function）。

generator_demo.py · 两种写法的对比

# ---------- 普通函数：返回完整列表，占内存 ----------
def squares_list(n: int) -> list[int]:
    """返回前 n 个平方组成的列表。"""
    return [x * x for x in range(n)]


# ---------- 生成器函数：产出器，每次只产出一个 ----------
def squares_gen(n: int):
    """逐个产出前 n 个平方，内存占用 O(1)。"""
    for x in range(n):
        yield x * x


# 对比内存与用法
a = squares_list(1_000_000)
print(type(a))                 # <class 'list'>
print(sys.getsizeof(a) // 1024, "KB")   # 约 8000 KB

b = squares_gen(1_000_000)
print(type(b))                 # <class 'generator'>
print(sys.getsizeof(b), "bytes")        # 约 200 bytes

# 生成器只能遍历一次
print(next(b))                 # 0
print(next(b))                 # 1
print(list(b)[:5])             # [4, 9, 16, 25, 36] —— 只剩下还没取过的部分

# 对比：list(生成器) 也可以，但会一次性物化
print(len(list(squares_gen(5))))   # 5

yield 的运行机制：暂停与恢复

理解 yield 的关键，是把它理解为「保存现场并返回」。每次 next()，生成器就执行到下一个 yield，把值交出去，同时冻结当前所有局部变量；下次 next()，从冻结处继续。

def counter():
    """演示 yield 如何保存与恢复现场。"""
    x = 10
    yield x            # 交出 10，此时 x=10 被保存
    x = x + 5
    yield x            # 交出 15，此时 x=15 被保存
    x = x * 2
    yield x            # 交出 30

g = counter()
print(next(g))   # 10
print(next(g))   # 15
print(next(g))   # 30
try:
    next(g)     # 已经没有更多 yield
except StopIteration:
    print("生成器已耗尽")

双向通信：send / throw / close

yield 表达式的值，就是下次调用 gen.send(v) 时传入的 v。这让生成器变成了可以双向对话的协程式对象。

echo_server.py · 用 send 实现「回声服务」

def echo_server():
    """一个回声服务生成器：收到什么就原样送回什么，收到 'quit' 就结束。"""
    while True:
        try:
            received = yield          # 第一次进来时接收 send(None)
        except GeneratorExit:
            print("[server] 连接被关闭")
            raise                     # 必须重新抛出 GeneratorExit
        if received == "quit":
            print("[server] 收到退出指令，结束服务")
            break
        print(f"[server] 收到: {received}")
        yield f"[echo] {received}"


srv = echo_server()

# 启动生成器，让它执行到第一个 yield
print(next(srv))            # None（第一次 send 必须是 None 或不传）

print(srv.send("你好"))      # [echo] 你好
print(srv.send("python"))   # [echo] python
print(srv.send("quit"))     # [server] 收到退出指令，结束服务

# close() 会向生成器抛 GeneratorExit，用于清理
g = echo_server()
next(g)
g.close()

生成器处理大文件的威力

log_analyzer.py · 用生成器流式分析日志

from collections import Counter
from pathlib import Path
from typing import Iterator


def read_lines(path: Path) -> Iterator[str]:
    """逐行读取文件，调用时才逐行进入内存。"""
    with open(path, encoding="utf-8", errors="replace") as f:
        for line in f:
            yield line.rstrip("\n")


def parse_errors(lines) -> Iterator[tuple[int, str]]:
    """从日志行中挑出 ERROR 级别的行，产出 (行号, 内容)。"""
    for no, line in enumerate(lines, start=1):
        if "ERROR" in line:
            yield no, line


def summarize(path: Path, top: int = 5) -> dict:
    """统计错误数与出现最多的错误类型。"""
    counter = Counter()
    total = 0
    for _, line in parse_errors(read_lines(path)):
        total += 1
        # 取 ERROR 后面第一个词作为错误类型
        counter[line.split("ERROR", 1)[1].strip().split()[0]] += 1
    return {"total_errors": total, "top": counter.most_common(top)}


# 演示：生成器可以像管道一样层层传递
def demo():
    import tempfile

    with tempfile.NamedTemporaryFile("w", suffix=".log", delete=False,
                                     encoding="utf-8") as f:
        for i in range(1, 6):
            f.write(f"2024-01-01 INFO 启动 {i}\n")
        f.write("2024-01-01 ERROR DatabaseTimeout 数据库超时\n")
        for i in range(6, 10):
            f.write(f"2024-01-01 INFO 请求 {i}\n")
        f.write("2024-01-01 ERROR DiskFull 磁盘已满\n")
        f.write("2024-01-01 ERROR DatabaseTimeout 数据库超时\n")
        name = f.name

    result = summarize(Path(name))
    print("错误总数 :", result["total_errors"])          # 3
    print("Top 错误 :", result["top"][0])                # ('DatabaseTimeout', 2)
    Path(name).unlink()


if __name__ == "__main__":
    demo()

yield from：委托给另一个生成器

手写 for x in sub(): yield x 冗长且低效（多了一层函数调用）。Python 3.3+ 提供 yield from（PEP 380）做委托。

# 对比：手写委托 vs yield from
def sub_gen():
    yield 1
    yield 2
    yield 3


def manual_delegate():
    for x in sub_gen():
        yield x * 10          # 还能在中间加工


def from_delegate():
    result = yield from sub_gen()   # 直接委托，无加工
    print("子生成器返回值:", result)
    yield 4


print(list(manual_delegate()))   # [10, 20, 30]
print(list(from_delegate()))     # [1, 2, 3, 4]（并打印 子生成器返回值: None）

# ---------- 真实场景：递归目录遍历 ----------
import os
from pathlib import Path


def walk_files(root: Path, suffix: str):
    """递归找出目录下所有指定后缀的文件（用 yield from 委托）。"""
    for entry in os.scandir(root):
        if entry.is_dir():
            yield from walk_files(Path(entry.path), suffix)
        elif entry.name.endswith(suffix):
            yield Path(entry.path)


def walk_files_manual(root: Path, suffix: str):
    """同样的功能，手写版本。"""
    for entry in os.scandir(root):
        if entry.is_dir():
            for f in walk_files_manual(Path(entry.path), suffix):
                yield f
        elif entry.name.endswith(suffix):
            yield Path(entry.path)


# 测试
import tempfile
with tempfile.TemporaryDirectory() as d:
    p = Path(d)
    (p / "a.txt").write_text("x")
    (p / "b.md").write_text("y")
    (p / "sub").mkdir()
    (p / "sub" / "c.txt").write_text("z")
    print(sorted(f.name for f in walk_files(p, ".txt")))   # ['a.txt', 'c.txt']
    print(sorted(f.name for f in walk_files_manual(p, ".txt")))

del_eggs.py · yield from 的隐藏能力：把迭代器的返回值传出来

def sub():
    """子生成器：return 的值会成为外层 yield from 表达式的值。"""
    yield 10
    yield 20
    return "子生成器完成啦"


def main():
    received = yield from sub()
    print("收到子生成器的返回值:", received)
    yield received.upper()


print(list(main()))
# 输出：
# 收到子生成器的返回值: 子生成器完成啦
# ['子生成器完成啦']

生成器表达式与惰性求值

如果生成器函数只用了一个 yield，就可以写成「生成器表达式（generator expression）」——一个用圆括号包裹的推导式。

# 列表推导式 vs 生成器表达式的语法差异只有一个字符：括号
squares_list = [x * x for x in range(1_000_000)]   # 方括号：立即物化
squares_gen  = (x * x for x in range(1_000_000))   # 圆括号：惰性

# 类型确认
print(type(squares_list))   # <class 'list'>
print(type(squares_gen))    # <class 'generator'>

# 内存对比
import sys
from sys import getsizeof
print(getsizeof(squares_list) // 1024, "KB")   # 约 8000 KB
print(getsizeof(squares_gen), "bytes")         # 约 200 bytes

# 惰性的另一个含义：外层迭代不执行时，内层根本不会算
data = [1, 2, 3, 4, 5]
gen = (x for x in data if x > 100)
print(list(gen))            # []  —— 判断根本没发生多少次
print(data)                 # 原列表没被修改

# 传给 sum / max / any / all：它们会自动消费生成器
print(sum(x * x for x in range(10)))         # 285
print(max((len(w) for w in ["a", "bb", "ccc"])))   # 3
print(any(x < 0 for x in [1, 2, -1]))      # True
print(all(x > 0 for x in [1, 2, 3]))       # True

# 嵌套推导式（列表可以嵌套，生成器表达式可以带条件过滤）
pairs = [(i, j) for i in range(3) for j in range(i, 3)]
print(pairs)   # [(0, 0), (0, 1), (0, 2), (1, 1), (1, 2), (2, 2)]
evens = (n for n in range(20) if n % 2 == 0)
print(list(evens))   # [0, 2, 4, ..., 18]

装饰器：一个函数改造另一个函数

装饰器（decorator）是 Python 里最常被滥用也最强大的工具。它本质上就是一个「接收函数、返回新函数」的高阶函数。

不用语法糖，先看本质

def my_decorator(func):
    """最基础的装饰器：包装 func，返回一个新函数。"""
    def wrapper(*args, **kwargs):
        print(f"[前] 调用 {func.__name__}，参数={args} {kwargs}")
        result = func(*args, **kwargs)
        print(f"[后] {func.__name__} 返回 {result}")
        return result
    return wrapper


def greet(name: str) -> str:
    """原始函数。"""
    return f"你好，{name}!"


# ===== 不用 @ 语法糖 =====
wrapped = my_decorator(greet)
print(wrapped("世界"))
# [前] 调用 greet，参数=('世界',) {}
# 你好，世界!
# [后] greet 返回 你好，世界!

# ===== 用 @ 语法糖，完全等价 =====
@my_decorator
def greet(name: str) -> str:
    return f"你好，{name}!"

print(greet("Python"))
# [前] 调用 greet，参数=('Python',) {}
# 你好，Python!
# [后] greet 返回 你好，Python!

元信息丢失问题

def naive_decorator(func):
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)
    return wrapper


@naive_decorator
def calculate(a: int, b: int) -> int:
    """计算两数之和。"""
    return a + b


print(calculate.__name__)       # wrapper      —— 原始名字丢了
print(calculate.__doc__)        # None         —— 文档字符串丢了
print(calculate.__module__)     # __main__     —— 模块信息也变了
import inspect
print(inspect.signature(calculate))   # (*args, **kwargs) —— 签名毁了
help(calculate)                 # 帮不上任何忙

解决办法是 functools.wraps：

from functools import wraps


def good_decorator(func):
    @wraps(func)                 # 把 func 的元信息复制到 wrapper
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)
    return wrapper


@good_decorator
def calculate(a: int, b: int) -> int:
    """计算两数之和。"""
    return a + b


print(calculate.__name__)       # calculate
print(calculate.__doc__)        # 计算两数之和。
print(calculate.__wrapped__)    # 指向原始函数，可以再取回来
import inspect
print(inspect.signature(calculate))   # (a: int, b: int) -> int
print(inspect.unwrap(calculate) is calculate.__wrapped__)   # True

保留原函数引用

from functools import wraps


def timer(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        import time
        t0 = time.perf_counter()
        result = func(*args, **kwargs)
        elapsed = time.perf_counter() - t0
        print(f"[计时] {func.__name__} 耗时 {elapsed:.6f}s")
        return result
    return wrapper


@timer
def slow_task():
    """一个假装很慢的任务。"""
    total = 0
    for i in range(200_000):
        total += i
    return total


print(slow_task())            # [计时] slow_task 耗时 0.0xxx s

# 通过 __wrapped__ 仍可访问原始实现（绕过装饰层）
print(slow_task.__wrapped__.__doc__)
# 要多次调用原始函数时用它；不要用它绕过必要的校验/日志

带参数的装饰器

如果装饰器需要「配置」（比如重试几次、超时多少秒），就要写成三明治结构：最外层接收配置，中间返回真正的装饰器。

retry.py · 重试装饰器

import time
from functools import wraps
from typing import Callable, TypeVar

F = TypeVar("F", bound=Callable[..., object])


def retry(times: int = 3, delay: float = 0.5, exceptions=(Exception,)):
    """失败自动重试的装饰器工厂。

    :param times: 最大尝试次数（含第一次）
    :param delay: 每次重试之间的等待秒数
    :param exceptions: 触发重试的异常类型元组
    """
    if times < 1:
        raise ValueError("times 必须 >= 1")

    def decorator(func: F) -> F:
        @wraps(func)
        def wrapper(*args, **kwargs):
            last_exc = None
            for attempt in range(1, times + 1):
                try:
                    return func(*args, **kwargs)
                except exceptions as exc:
                    last_exc = exc
                    if attempt == times:
                        break
                    print(f"[重试] {func.__name__} 第 {attempt} 次失败：{exc!r}，{delay}s 后再试")
                    time.sleep(delay)
            # 所有尝试都失败，抛出最后一次异常
            raise last_exc
        return wrapper
    return decorator


# ---------- 使用 ----------
import random

attempts = {"n": 0}


@retry(times=5, delay=0.05, exceptions=(ConnectionError,))
def flaky_api() -> str:
    """模拟一个时好时坏的接口。"""
    attempts["n"] += 1
    if attempts["n"] < 3:
        raise ConnectionError("连接被重置")
    return "成功返回数据"


print(flaky_api())        # 两次重试日志后输出「成功返回数据」
print("总尝试次数:", attempts["n"])   # 3


# 不带括号时无法传参——必须用括号
@retry()                   # 默认 times=3
def another() -> str:
    return "OK"


print(another())

类装饰器与有状态装饰器

装饰器不一定是函数，任何可调用对象都可以。类的实例（实现了 __call__）常常更适合有状态的装饰器。

class_decorators.py · 用类实现的装饰器

from functools import wraps
from collections import Counter


class CountCalls:
    """统计函数被调用次数的类装饰器（有状态）。"""

    def __init__(self, func):
        self.func = func
        self.counter = Counter()
        wraps(func)(self)      # 把 func 的元信息挂到这个实例上

    def __call__(self, *args, **kwargs):
        key = f"{self.func.__name__}:{args}:{sorted(kwargs.items())}"
        self.counter[key] += 1
        print(f"[计数] {self.func.__name__} 第 {self.counter[key]} 次调用")
        return self.func(*args, **kwargs)

    def stats(self) -> dict:
        return dict(self.counter)


@CountCalls
def greet(name: str) -> str:
    return f"你好，{name}!"


print(greet("A"))    # [计数] greet 第 1 次调用 / 你好，A!
print(greet("B"))    # ... 第 1 次调用
print(greet("A"))    # ... 第 2 次调用
print(greet.stats()) # {"greet:('A',):[]": 2, "greet:('B',):[]": 1}
print(greet.__name__)   # greet —— 实例伪装得像原函数


# ---------- 另一个例子：类装饰器修改类 ----------
import functools


def add_repr(cls):
    """给类自动补一个可读的 __repr__。"""
    @functools.wraps(cls)
    def wrapper(*args, **kwargs):
        obj = cls(*args, **kwargs)
        obj.info = f"{cls.__name__}(构造完成)"
        return obj
    return wrapper


@add_repr
class Service:
    def __init__(self, name: str) -> None:
        self.name = name


s = Service("订单服务")
print(s.info)   # Service(构造完成)
print(s.name)   # 订单服务

装饰器叠加：从下往上应用

import functools
import time

def log_calls(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print(f"[log] -> {func.__name__}")
        return func(*args, **kwargs)
    return wrapper


def count_time(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        t0 = time.perf_counter()
        result = func(*args, **kwargs)
        print(f"[time] {func.__name__} 用时 {time.perf_counter() - t0:.6f}s")
        return result
    return wrapper


@log_calls
@count_time
def work(n: int) -> int:
    """被两层装饰的函数。"""
    return sum(range(n))


work(100_000)
# [log] -> work
# [time] work 用时 0.00xxxs
#
# 等价于 work = log_calls(count_time(work))
# 所以 log_calls 在最外层：先被 log 看到，再进入计时，再进 work

带参数的类装饰器

import functools


def register(name: str):
    """把函数注册进全局注册表。"""

    def decorator(func):
        REGISTRY[name] = func

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)
        return wrapper
    return decorator


REGISTRY: dict[str, object] = {}


@register("double")
def _double(x: int) -> int:
    return x * 2


@register("triple")
def _triple(x: int) -> int:
    return x * 3


for cmd, fn in REGISTRY.items():
    print(cmd, "->", fn(5))     # double -> 10 / triple -> 15

实用装饰器工具箱

下面这些装饰器在任何真实项目里都会反复用到，值得逐个理解并抄进自己的 utils/decorators.py。

1. 计时装饰器

import time
import functools


def timeit(func):
    """打印函数执行耗时的装饰器。"""

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        t0 = time.perf_counter()
        try:
            return func(*args, **kwargs)
        finally:
            dt = time.perf_counter() - t0
            print(f"{func.__name__} 耗时 {dt:.6f}s")
    return wrapper


@timeit
def compute(n: int) -> int:
    return sum(range(n))


compute(1_000_000)      # compute 耗时 0.01xxx s

2. 内存缓存（memoize）

import functools


@functools.cache                       # Python 3.9+ 内置版 lru_cache(maxsize=None)
def fib(n: int) -> int:
    """朴素递归斐波那契，加缓存后从指数级降到 O(n)。"""
    return n if n < 2 else fib(n - 1) + fib(n - 2)


print(fib(50))        # 12586269025（瞬间完成）
print(fib.cache_info())   # CacheInfo(hits=48, misses=51, maxsize=None, currsize=51)

# 限制大小的版本：适合不希望无限增长的场景
@functools.lru_cache(maxsize=128)
def parse_line(line: str) -> tuple:
    return tuple(line.split(","))


print(parse_line("a,b,c"))
print(parse_line.cache_info())
parse_line.cache_clear()      # 手动清空

3. 参数校验装饰器

import functools
import inspect


def require_types(**expected):
    """按参数名校验类型，不满足就抛 TypeError。

    用法：@require_types(x=int, y=(int, float))
    """

    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            bound = inspect.signature(func).bind(*args, **kwargs)
            for name, want in expected.items():
                if name not in bound.arguments:
                    continue
                value = bound.arguments[name]
                if not isinstance(value, want):
                    raise TypeError(
                        f"{func.__name__} 的参数 {name} 期望 {want}，"
                        f"实际得到 {type(value).__name__}"
                    )
            return func(*args, **kwargs)
        return wrapper
    return decorator


@require_types(a=int, b=(int, float))
def divide(a, b):
    """两数相除。"""
    return a / b


print(divide(10, 2))            # 5.0
try:
    divide("x", 2)
except TypeError as e:
    print("捕获到:", e)          # 捕获到: divide 的参数 a 期望 (<class 'int'>)，实际得到 str

4. 权限检查装饰器

import functools

CURRENT_USER = {"name": "alice", "roles": {"admin"}}


def require_role(*roles: str):
    """要求当前用户拥有其中任一角色才能调用。"""

    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            user_roles = CURRENT_USER["roles"]
            if not user_roles.intersection(roles):
                raise PermissionError(
                    f"用户 {CURRENT_USER['name']} 缺少角色 {roles}"
                )
            return func(*args, **kwargs)
        return wrapper
    return decorator


@require_role("admin", "ops")
def delete_all() -> str:
    """危险操作。"""
    return "已删除全部数据"


print(delete_all())      # 已删除全部数据

CURRENT_USER["roles"] = {"guest"}
try:
    delete_all()
except PermissionError as e:
    print("拒绝:", e)    # 拒绝: 用户 alice 缺少角色 ('admin', 'ops')

5. 弃用警告装饰器

import functools
import warnings


def deprecated(reason: str = "", version: str = ""):
    """调用被弃用的函数时发出 DeprecationWarning。"""

    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            msg = f"{func.__name__} 已弃用"
            if version:
                msg += f"（将在 {version} 移除）"
            if reason:
                msg += f"：{reason}"
            warnings.warn(msg, DeprecationWarning, stacklevel=2)
            return func(*args, **kwargs)
        return wrapper
    return decorator


@deprecated(reason="请改用 new_api()", version="3.0")
def old_api(x: int) -> int:
    return x * 2


import warnings
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    print(old_api(3))          # 6
    print(caught[0].message)   # old_api 已弃用（将在 3.0 移除）：请改用 new_api()

上下文管理器：with 的实现

第 8 章我们已经用过 with open(...) as f:。它背后同样是「函数即对象」的机制。

上下文管理器协议

一个对象要能用于 with，需要实现：

1. __enter__()：进入时调用，返回值赋给 as 后面的名字。
2. __exit__(exc_type, exc_value, traceback)：退出时一定被调用（正常或异常）。返回真值表示「吞掉异常」，否则异常继续向上抛。

context_manager_demo.py · 手写一个上下文管理器类

class Timer:
    """测量 with 块内代码耗时的上下文管理器。"""

    def __init__(self, label: str = "代码块") -> None:
        self.label = label
        self.start = 0.0
        self.elapsed = 0.0

    def __enter__(self) -> "Timer":
        import time
        self.start = time.perf_counter()
        print(f"[{self.label}] 开始")
        return self            # 返回值会绑定给 as 后的名字

    def __exit__(self, exc_type, exc_value, tb) -> bool:
        import time
        self.elapsed = time.perf_counter() - self.start
        if exc_type is None:
            print(f"[{self.label}] 正常结束，耗时 {self.elapsed:.6f}s")
        else:
            print(f"[{self.label}] 因 {exc_type.__name__} 结束，耗时 {self.elapsed:.6f}s")
        return False            # False = 不吞异常


with Timer("数据处理") as t:
    total = sum(range(500_000))
print("计算结果:", total)
print("记录耗时:", t.elapsed)      # 通过 as 拿到 __enter__ 的返回值


# 异常也会触发 __exit__
try:
    with Timer("危险操作") as t2:
        raise RuntimeError("出错了")
except RuntimeError as e:
    print("异常已传播:", e)

@contextmanager：函数写法

大多数时候用装饰器更方便。注意结构：yield 之前是进入逻辑，之后是退出逻辑。

from contextlib import contextmanager
import time


@contextmanager
def open_file(path: str, mode: str = "r"):
    """打开文件的上下文管理器（即使出错也保证关闭）。"""
    print(f"[open] {path}")
    f = open(path, mode, encoding="utf-8")
    try:
        yield f                       # 这里把 f 交给 with 的调用方
    except Exception as exc:
        print(f"[error] 读取 {path} 出错：{exc!r}")
        raise                         # 决定是否向上抛
    finally:
        f.close()                     # 无论正常还是异常都会执行
        print(f"[close] {path}")


# 测试
import tempfile, os
name = tempfile.mktemp(suffix=".txt")
with open(name, "w", encoding="utf-8") as f:
    f.write("第一行\n第二行\n")

with open_file(name) as f:
    print("内容:", repr(f.read()))

os.unlink(name)


# 另一个：临时目录
@contextmanager
def temp_dir():
    """自动创建并清理临时目录。"""
    import tempfile, shutil
    d = tempfile.mkdtemp(prefix="demo_")
    try:
        yield d
    finally:
        shutil.rmtree(d, ignore_errors=True)


with temp_dir() as d:
    print("临时目录:", d)
    open(os.path.join(d, "x.txt"), "w").write("hi")
print("离开后目录已清理:", not os.path.exists(d))

contextlib 里的现成工具

import io
import os
import tempfile
from contextlib import ExitStack, redirect_stdout, suppress


# 1. 忽略异常
with suppress(FileNotFoundError):
    os.remove("/绝对不存在的路径.txt")
print("删除失败被安全忽略，程序继续")


# 2. 重定向 stdout：常用于测试与日志捕获
buf = io.StringIO()
with redirect_stdout(buf):
    print("这条不会打印到终端")
captured = buf.getvalue().strip()
print("已捕获:", captured)


# 3. ExitStack：数量不确定的资源
with ExitStack() as stack:
    files = []
    for name in ["a.txt", "b.txt", "c.txt"]:
        f = stack.enter_context(open(name, "w", encoding="utf-8"))
        f.write(f"内容 {name}")
        files.append(f)
    print("已同时打开", len(files), "个文件")
# 离开 with 时按相反顺序自动关闭全部文件
print("全部已关闭:", all(f.closed for f in files))
for name in ["a.txt", "b.txt", "c.txt"]:
    os.remove(name)

实际例子：数据库事务

transaction.py · 用上下文管理器实现事务

from contextlib import contextmanager


class FakeDB:
    """模拟数据库连接，用来演示事务语义。"""

    def __init__(self) -> None:
        self.in_tx = False
        self.undo_log: list[str] = []

    def begin(self) -> None:
        self.in_tx = True
        self.undo_log.clear()
        print("[db] BEGIN")

    def execute(self, sql: str) -> None:
        if not self.in_tx:
            raise RuntimeError("不在事务中")
        if sql.startswith("UPDATE"):
            self.undo_log.append(sql)
        print("[db]", sql)

    def commit(self) -> None:
        self.in_tx = False
        print("[db] COMMIT")


@contextmanager
def transaction(db: FakeDB):
    """事务上下文：正常则提交，异常则回滚。"""
    db.begin()
    try:
        yield db
    except Exception:
        print("[db] ROLLBACK:", db.undo_log or "无变更")
        db.in_tx = False
        raise
    else:
        db.commit()
    finally:
        db.undo_log.clear()


# 成功路径
db = FakeDB()
with transaction(db) as conn:
    conn.execute("UPDATE account SET balance = 0")
    conn.execute("UPDATE account SET balance = 900")
print("事务后状态:", db.in_tx)


# 失败路径：自动回滚，异常继续抛出
db2 = FakeDB()
try:
    with transaction(db2) as conn:
        conn.execute("UPDATE account SET balance = 0")
        raise ValueError("业务校验失败")
except ValueError as e:
    print("业务异常:", e)
print("已回滚:", not db2.in_tx)

三个概念的统一视角

迭代器、生成器、装饰器看似无关，其实都建立在同一个基础上：「函数是对象」。

选择决策表

最佳实践与常见错误

必须遵守的规则

1. 装饰器永远加 @wraps(func)。
2. 装饰器保持透明：返回原函数的返回值，不要擅自吞掉参数或改变语义。
3. 清理代码放 finally：装饰器与上下文管理器的清理逻辑都要保证异常路径下也执行。
4. 不要在生成器里写 return 值 给循环用：普通 for 拿不到，只有 yield from 能接住。
5. 明确迭代器的一次性：需要复用就 list(gen) 或重新生成。

三个典型错误代码

import functools

# ❌ 错误一：返回列表而非函数，调用会失败
def bad_decorator(func):
    return [func, "extra"]


# ❌ 错误二：把 *args 直接传给另一个普通函数（需要 kwargs 透传时漏掉）
def bad2(func):
    def wrapper(x):
        return func(x)        # ❌ 少传了 **kwargs
    return wrapper


# ❌ 错误三：循环闭包：装饰器/lambda 捕获的是变量而非值
funcs = [lambda: i for i in range(3)]
print([f() for f in funcs])       # [2, 2, 2] —— 全是 2

# ✅ 修正：默认参数绑定当前值
funcs = [lambda i=i: i for i in range(3)]
print([f() for f in funcs])       # [0, 1, 2]


# ❌ 错误四：生成器只被部分消费就丢弃
def read_all():
    for line in open("big.log", encoding="utf-8"):
        yield line

gen = read_all()
first = next(gen)
# 如果这里就丢掉 gen，生成器不会自动关闭；用 with 或显式 close
gen.close()


# ✅ 修正：正确的带参装饰器模板
def logging_with_level(level: str = "INFO"):
    import logging
    logger = logging.getLogger(__name__)

    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            logger.log(logging.INFO, "%s 正在执行，参数=%s", func.__name__, args)
            try:
                return func(*args, **kwargs)
            finally:
                logger.log(logging.INFO, "%s 已结束", func.__name__)
        return wrapper
    return decorator


@logging_with_level("INFO")
def business_op(x: int) -> int:
    return x + 1


print(business_op(1))    # 2（日志未配置 handler 时不显示）

自检清单

• □ 能说出 for x in obj 的完整展开过程。
• □ 能手写一个实现了 __iter__/__next__ 的迭代器类。
• □ 写涉及大量数据的函数时优先考虑生成器，并解释内存差异。
• □ 所有自定义装饰器都加了 @wraps。
• □ 带参装饰器能正确写出「工厂 → 装饰器 → 包装器」三层。
• □ 能用 @contextmanager 封装资源管理，且清理逻辑在 finally 中。

本章小结

• 迭代协议：__iter__ 返回迭代器，__next__ 产出元素，StopIteration 结束；for 就是这三步的语法糖。
• 可迭代 ≠ 迭代器：容器可重复遍历，迭代器（文件、生成器、map/zip）一次性消耗。
• 生成器：含 yield 的函数返回生成器对象，内存 O(1)，可暂停/恢复；send/throw/close 支持双向通信。
• yield from：委托子生成器并接收其 return 值，比手写 for-yield 快且清晰。
• 生成器表达式（圆括号）用于惰性求值；用一次选生成器，要复用选列表。
• 装饰器是「函数 → 函数」的高阶函数；@dec 等价于 func = dec(func)；必须加 functools.wraps。
• 带参装饰器用三层结构：retry(...) → decorator(func) → wrapper(...)；多层装饰自下而上应用。
• 类装饰器通过 __call__ 实现，适合有状态的装饰（如调用计数、注册表）。
• 上下文管理器实现 __enter__/__exit__；函数版用 @contextmanager，结构是「yield 前进入、yield 后清理（finally）」。
• 标准工具：functools.cache/lru_cache/wraps/partial、contextlib.suppress/ExitStack/redirect_stdout。
• 典型 bug：循环闭包 late binding、忘 @wraps、@contextmanager 中 yield 次数不对、生成器被部分消费。

练习题

1. 手写迭代器：实现类 Range3，行为与内置 range 一致（支持 start/stop/step，支持负步长，支持 len() 和 __contains__），并能被 list() 收集。

   参考答案要点实现 __iter__ 返回自身、__next__ 判越界后抛 StopIteration；__len__ 返回 max(0, (stop-start+step-1)//step)（step>0 情形，负步长类推）；__contains__ 可用 self.start <= v < self.stop 快速判断。写好后用 assert list(Range3(0, 10, 3)) == [0, 3, 6, 9] 验证。

2. 内存对比实验：写一段脚本，分别用列表推导式与生成器表达式生成 1000 万个偶数的平方，打印两者的 sys.getsizeof 与 time 耗时，并解释差异。

   参考答案要点a = [i*i for i in range(10_000_000)] 与 b = (i*i for i in range(10_000_000))。列表约 80 MB 且构建耗时长；生成器对象约 200 字节，若 sum(b) 消费则峰值内存极低但耗时会略高（逐个 yield 有开销）。要点：列表「快占内存」，生成器「省内存省时间地开始」。

3. 生成器管道：用三个生成器函数实现整数流水处理：gen_source() 产出 1–100，gen_filter(numbers) 只留偶数，gen_map(numbers, fn) 把每个数平方，consume(gen) 求总和。禁止使用 list 中转。

   参考答案

   def gen_source():
       for i in range(1, 101):
           yield i
   
   def gen_filter(numbers):
       for n in numbers:
           if n % 2 == 0:
               yield n
   
   def gen_map(numbers, fn):
       for n in numbers:
           yield fn(n)
   
   def consume(gen):
       total = 0
       for v in gen:
           total += v
       return total
   
   print(consume(gen_map(gen_filter(gen_source()), lambda x: x * x)))  # 2^2+4^2+...+100^2
   

   要点：每个函数只关心「上一步给我一个可迭代对象」，这就是惰性管道的解耦威力。

4. 带参装饰器：实现 throttle(interval)，被装饰函数两次调用的间隔小于 interval 秒时直接返回 None（限流），用 @wraps 保留元信息。

   参考答案要点用 time.monotonic() 记录上次调用时间（不要用 time.time()，它会被系统时钟调整影响）；闭包里保存 last = [0.0]；若 now - last[0] < interval 则返回 None，否则更新并调用原函数。注意 @wraps 不能省。

5. 装饰器叠加顺序：有 @a、@b、@c 三个装饰器，写出完整展开形式，并说明函数调用时打印顺序（每个装饰器在调用前和调用后各打印一次）。

   参考答案展开形式：f = a(b(c(f)))。调用顺序：进入时 a前 → b前 → c前 → f → c后 → b后 → a后（因为 a 在最外层）。定义时应用顺序是自下而上（先 c 后 b 后 a），执行时进入自上而下、退出自下而上。这是理解多层装饰日志顺序的关键。

6. 类装饰器：实现类装饰器 Count，为被装饰函数自动添加 count 属性（记录调用次数）与 reset() 方法，装饰后仍能正常调用原函数。

   参考答案要点

   import functools
   
   class Count:
       def __init__(self, func):
           self.func = func
           self.count = 0
           functools.wraps(func)(self)
   
       def __call__(self, *a, **kw):
           self.count += 1
           return self.func(*a, **kw)
   
       def reset(self):
           self.count = 0
   

   装饰后 f.count 与 f.reset() 都能访问（因为实例把属性暴露给了包装后的名字）。注意 wraps(func)(self) 会尝试设置 self.__name__ 等属性，若原函数有 __dict__ 需注意冲突。

7. 上下文管理器：用 @contextmanager 写 temporary_file(content)，自动创建临时文件并写入内容，退出时删除。演示异常路径下文件同样被删除。

   参考答案

   import os, tempfile
   from contextlib import contextmanager
   
   @contextmanager
   def temporary_file(content: str):
       f = tempfile.NamedTemporaryFile("w", suffix=".txt",
                                        delete=False, encoding="utf-8")
       path = f.name
       try:
           f.write(content)
           f.flush()
           yield path
       finally:
           if not f.closed:
               f.close()
           os.unlink(path)
   

   测试：with temporary_file("hello") as p: print(open(p).read())；再在 with 块内 raise ValueError，块外用 os.path.exists(p) 验证为 False。关键是删除写在 finally 里。

8. 缓存装饰器陷阱：@functools.cache 装饰的函数返回 list，调用方 append 了一下，第二次调用结果被污染。解释原因并给出两种修复方案。

   参考答案原因：cache 按参数缓存返回值，而 list 是可变对象，缓存的是同一个引用。修复：① 只对不可变返回值使用缓存（改成返回 tuple）；② 缓存内部序列、每次返回副本——写 @cache 的私有函数 _cached(*args) -> tuple，对外函数返回 list(_cached(*args))。还可加 info.clear() 提供手动清缓存的入口。

9. 生成器 close 与资源释放：写一个生成器函数，函数体里用 with open(...) 打开文件。解释：如果只消费前 5 行就丢弃生成器，文件何时被关闭？为什么？

   参考答案要点生成器被丢弃时引用计数归零，会触发其 close()，进而在挂起位置抛出 GeneratorExit；由于 with 块在生成器帧的栈上，with 的 __exit__ 会被调用，文件随即关闭。但更稳妥的写法是把生成器放在 with contextlib.closing(gen): 中，或用 try/finally 在生成器内部保证释放——依赖析构时机在有循环引用时并不可靠。

10. 综合工具库：写一个 utils/decorators.py，包含四个装饰器：log_calls（打印进入/退出）、retry（可配次数与异常）、timed（打印耗时）、validate(**types)（类型校验），并写一个 demo 脚本把它们叠加使用，验证元信息与行为都正确。

    参考答案要点每个装饰器都以 @functools.wraps(func) 开头、参数用 *args, **kwargs 透传；timed 的计时与打印放 finally 以覆盖异常路径；retry 只捕获指定异常且最后一次不再重试；validate 用 inspect.signature(func).bind(*args, **kwargs) 绑定参数再逐个 isinstance 校验。叠加顺序示例：@log_calls 在 @retry 之外，这样重试过程也只记一条日志。