Python 上下文管理器:with 语句的真正力量


with 语句在干什么?

with open("file.txt") as f:
    data = f.read()
# 文件一定会被关闭,即使中间报错

with 不是魔法,它只做两件事:

  1. 进入时调用 __enter__() — 获取资源
  2. 退出时调用 __exit__() — 释放资源(无论是否异常)

自己写一个上下文管理器

class Timer:
    def __enter__(self):
        import time
        self.start = time.perf_counter()
        return self  # as 后面的变量

    def __exit__(self, exc_type, exc_val, exc_tb):
        import time
        self.elapsed = time.perf_counter() - self.start
        print(f"耗时 {self.elapsed:.4f}s")
        return False  # 不吞异常

with Timer() as t:
    sum(range(1_000_000))
# 耗时 0.0234s

__exit__ 的三个参数:

  • exc_type — 异常类型(无异常时为 None)
  • exc_val — 异常实例
  • exc_tb — traceback
  • 返回 True = 吞掉异常,返回 False = 继续抛出

contextlib — 用生成器写上下文管理器

不想写类?用 @contextmanager 更简洁:

from contextlib import contextmanager
import time

@contextmanager
def timer(label=""):
    start = time.perf_counter()
    try:
        yield  # 这里暂停,执行 with 块的代码
    finally:
        elapsed = time.perf_counter() - start
        print(f"{label} 耗时 {elapsed:.4f}s")

with timer("计算"):
    sum(range(1_000_000))
# 计算 耗时 0.0234s

yield 之前 = __enter__yield 之后 = __exit__


10 个实战场景

1. 数据库连接

@contextmanager
def db_connection(url):
    conn = create_connection(url)
    try:
        yield conn
        conn.commit()  # 正常退出则提交
    except Exception:
        conn.rollback()  # 异常则回滚
        raise
    finally:
        conn.close()  # 一定关闭

with db_connection("postgres://...") as conn:
    conn.execute("INSERT INTO ...")

2. 临时改变工作目录

import os
from contextlib import contextmanager

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

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

3. 临时修改环境变量

@contextmanager
def env_var(key, value):
    old = os.environ.get(key)
    os.environ[key] = value
    try:
        yield
    finally:
        if old is None:
            del os.environ[key]
        else:
            os.environ[key] = old

with env_var("DEBUG", "1"):
    assert os.environ["DEBUG"] == "1"
# 自动恢复

4. 线程锁

import threading

lock = threading.Lock()

# with 自动获取和释放锁
with lock:
    # 线程安全的代码
    shared_resource.update()

5. 临时重定向输出

from contextlib import redirect_stdout
import io

buffer = io.StringIO()
with redirect_stdout(buffer):
    print("captured!")

print(buffer.getvalue())  # "captured!\n"

6. 忽略特定异常

from contextlib import suppress

# 不 care 文件不存在
with suppress(FileNotFoundError):
    os.remove("maybe_exists.txt")

# 等价于:
# try:
#     os.remove("maybe_exists.txt")
# except FileNotFoundError:
#     pass

7. 临时文件/目录

import tempfile

with tempfile.NamedTemporaryFile(suffix=".csv") as f:
    f.write(b"data")
    f.flush()
    process(f.name)
# 自动删除

with tempfile.TemporaryDirectory() as tmpdir:
    # 在临时目录里工作
    save_files(tmpdir)
# 自动清理整个目录

8. 多资源同时管理

# Python 3.1+
with open("input.txt") as fin, open("output.txt", "w") as fout:
    for line in fin:
        fout.write(line.upper())

# Python 3.10+ 可以用括号换行
with (
    open("a.txt") as a,
    open("b.txt") as b,
    open("c.txt", "w") as c,
):
    ...

9. 异步上下文管理器

class AsyncDB:
    async def __aenter__(self):
        self.conn = await create_async_connection()
        return self.conn

    async def __aexit__(self, *exc):
        await self.conn.close()

async with AsyncDB() as conn:
    await conn.execute("SELECT 1")

10. ExitStack — 动态管理多个上下文

from contextlib import ExitStack

def process_files(file_list):
    with ExitStack() as stack:
        files = [stack.enter_context(open(f)) for f in file_list]
        # 所有文件都会在退出时关闭
        for f in files:
            process(f)

知识卡片 📌

上下文管理器 = 资源的自动获取与释放

类写法:
  __enter__()  →  获取资源
  __exit__()   →  释放资源

生成器写法(推荐):
  @contextmanager
  yield 之前 = enter
  yield 之后 = exit

常用内置:
  suppress()        — 忽略异常
  redirect_stdout() — 重定向输出
  ExitStack()       — 动态管理多个上下文
  tempfile.*        — 临时文件/目录

记住:
  with 不只是 try/finally 的语法糖
  它是"资源生命周期管理"的 Python 范式