上下文管理器

为什么需要 with

很多资源用完必须关：文件、数据库连接、锁、网络连接。

# 不用 with：手动关，容易忘
f = open("data.txt")
content = f.read()
f.close()           # 出异常就漏了

# 用 with：自动关
with open("data.txt") as f:
    content = f.read()
# 出 with 块 → 自动 close()，即使报错也会关

自己写一个上下文管理器（类版本）

实现 __enter__ 和 __exit__：

class Timer:
    def __enter__(self):
        import time
        self.start = time.time()
        return self                    # as 后面拿到的就是它

    def __exit__(self, exc_type, exc_val, tb):
        import time
        cost = time.time() - self.start
        print(f"耗时 {cost*1000:.2f} ms")
        return False                   # False = 不吞掉异常


with Timer():
    sum(range(10_000_000))
# 耗时 195.32 ms

exit 三个参数

返回 True = 吞掉异常；返回 False（或不返回）= 继续抛出。

@contextmanager：用 yield 写更简单

类版本要写两个魔法方法。contextmanager 用 yield 三行搞定：

from contextlib import contextmanager
import time

@contextmanager
def timer():
    start = time.time()
    try:
        yield                          # yield 之前 = __enter__
    finally:
        cost = time.time() - start
        print(f"耗时 {cost*1000:.2f} ms")
                                       # yield 之后 = __exit__

with timer():
    sum(range(10_000_000))

90% 自定义上下文管理器都该用 @contextmanager 写。

临时改目录的小工具

@contextmanager
def cd(path):
    import os
    old = os.getcwd()
    os.chdir(path)
    try:
        yield
    finally:
        os.chdir(old)

with cd("/tmp"):
    print(os.getcwd())     # /tmp
print(os.getcwd())          # 回到原来的目录

临时设置环境变量

@contextmanager
def env(**kwargs):
    import os
    old = {k: os.environ.get(k) for k in kwargs}
    os.environ.update(kwargs)
    try:
        yield
    finally:
        for k, v in old.items():
            if v is None:
                os.environ.pop(k, None)
            else:
                os.environ[k] = v

with env(DEBUG="1", LOG_LEVEL="DEBUG"):
    run_tests()

多个上下文管理器

with open("a.txt") as fa, open("b.txt", "w") as fb:
    fb.write(fa.read())

或用 ExitStack 管理一组：

from contextlib import ExitStack

with ExitStack() as stack:
    files = [stack.enter_context(open(p)) for p in paths]
    # 出 with 自动全关

库里常见的上下文管理器

下一篇进入 OOP 进阶——类的继承与多态。