上下文管理器(Context Manager)是 Python 中用于资源管理的核心机制。无论你是在读写文件、获取锁、打开数据库连接还是管理网络会话——with 语句都能确保资源被正确释放,即使发生异常也不例外。

为什么需要上下文管理器

传统方式的痛点

来看一个常见的文件操作:

# 传统方式:容易忘记关闭文件
f = open("data.txt", "r")
content = f.read()
f.close()  # 如果前面的代码抛异常,这一行不会执行!


# 改进:用 try/finally 确保关闭
f = open("data.txt", "r")
try:
    content = f.read()
finally:
    f.close()  # 无论是否异常,都会执行

虽然 try/finally 能解决问题,但代码显得冗长。如果每次资源操作都要写 try/finally,不仅繁琐而且容易遗漏。

使用 with 语句

# with 语句 —— 简洁且安全
with open("data.txt", "r") as f:
    content = f.read()
# 离开 with 块后,文件自动关闭


# 等价于 try/finally,但更简洁

with 语句的工作原理

with 语句的本质是执行上下文管理器的协议:

with EXPRESSION as VARIABLE:
    BLOCK

等价于:

manager = EXPRESSION
value = manager.__enter__()
VARIABLE = value
try:
    BLOCK
except Exception as e:
    if not manager.__exit__(type(e), e, e.__traceback__):
        raise
else:
    manager.__exit__(None, None, None)

实现上下文管理器

方式一:类实现 __enter____exit__

class ManagedFile:
    def __init__(self, filename: str, mode: str = "r"):
        self.filename = filename
        self.mode = mode
        self.file = None

    def __enter__(self):
        print(f"打开文件: {self.filename}")
        self.file = open(self.filename, self.mode)
        return self.file  # 返回给 as 子句

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.file:
            self.file.close()
            print(f"关闭文件: {self.filename}")

        # 返回 True 会抑制异常,返回 False(默认)则异常继续传播
        if exc_type is not None:
            print(f"发生异常: {exc_val}")
        return False  # 不抑制异常


# 正常使用
with ManagedFile("test.txt", "w") as f:
    f.write("Hello, World!")
# 输出:
# 打开文件: test.txt
# 关闭文件: test.txt

# 异常时的使用
try:
    with ManagedFile("test.txt", "r") as f:
        content = f.read()
        raise ValueError("模拟错误")
except ValueError as e:
    print(f"捕获到异常: {e}")
# 输出:
# 打开文件: test.txt
# 发生异常: 模拟错误
# 关闭文件: test.txt
# 捕获到异常: 模拟错误

方式二:@contextmanager 装饰器

更简洁的方式是使用 contextlib 模块的 @contextmanager 装饰器:

from contextlib import contextmanager


@contextmanager
def managed_file(filename: str, mode: str = "r"):
    """使用生成器实现的上下文管理器"""
    print(f"打开文件: {filename}")
    file = open(filename, mode)
    try:
        yield file  # yield 之前的代码相当于 __enter__
    finally:
        file.close()
        print(f"关闭文件: {filename}")


with managed_file("test.txt", "w") as f:
    f.write("Hello from generator!")

@contextmanager 的异常处理

@contextmanager
def divider(a: float, b: float):
    """安全的除法上下文管理器"""
    print(f"准备计算: {a} / {b}")
    try:
        yield a / b
    except ZeroDivisionError as e:
        print(f"计算错误: {e}")
        yield float("inf")  # 返回无穷大作为替代值
    finally:
        print("计算完成")


with divider(10, 2) as result:
    print(f"结果: {result}")       # 结果: 5.0

with divider(10, 0) as result:
    print(f"结果: {result}")       # 结果: inf

contextlib 实用工具

Python 的 contextlib 模块提供了丰富的上下文管理器工具:

closing

确保对象被正确关闭(调用 close() 方法):

from contextlib import closing
import urllib.request


with closing(urllib.request.urlopen("https://httpbin.org/get")) as response:
    data = response.read()
    print(f"状态码: {response.status}")
# 离开 with 块后,response.close() 自动调用

suppress

忽略指定的异常:

from contextlib import suppress
import os


# 传统方式
try:
    os.remove("temp.txt")
except FileNotFoundError:
    pass

# 使用 suppress
with suppress(FileNotFoundError):
    os.remove("temp.txt")

# 可以同时忽略多种异常
with suppress(FileNotFoundError, PermissionError):
    os.remove("temp.txt")

nullcontext

需要一个上下文管理器但什么都不做时使用:

from contextlib import nullcontext
from typing import ContextManager


def process_data(data: str, use_lock: bool = False) -> None:
    """根据参数决定是否需要加锁"""
    lock = threading.Lock() if use_lock else nullcontext()

    with lock:  # 有锁时加锁,没有锁时 nullcontext 什么也不做
        print(f"Processing: {data}")

redirect_stdoutredirect_stderr

临时重定向标准输出/错误:

from contextlib import redirect_stdout, redirect_stderr
import io


def noisy_function():
    print("This goes to stdout")
    print("This goes to stderr", file=sys.stderr)


# 捕获输出到字符串
stdout_capture = io.StringIO()
stderr_capture = io.StringIO()

with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
    noisy_function()

print(f"捕获的 stdout: {stdout_capture.getvalue()}")
print(f"捕获的 stderr: {stderr_capture.getvalue()}")

# 恢复后,输出回到正常位置
noisy_function()

ExitStack

动态管理多个上下文管理器:

from contextlib import ExitStack


def open_many_files(filenames: list[str]) -> list:
    """打开多个文件,确保所有文件都能被正确关闭"""
    files = []
    with ExitStack() as stack:
        for filename in filenames:
            f = open(filename, "w")
            stack.callback(f.close)  # 注册清理回调
            files.append(f)
        # 返回前将文件列表与 ExitStack 关联
        return files  # 注意:这里文件会在 with 块退出时被关闭!


# 正确做法:让调用者也使用 with
def process_files(filenames: list[str]) -> None:
    with ExitStack() as stack:
        files = [
            stack.enter_context(open(fname))
            for fname in filenames
        ]
        # 在 with 块内安全使用所有文件
        for f, name in zip(files, filenames):
            f.write(f"Content for {name}")

嵌套上下文管理器

多个 with 的嵌套

# 嵌套写法
with open("input.txt") as infile:
    with open("output.txt", "w") as outfile:
        outfile.write(infile.read())

# 等效的扁平写法(Python 3.1+)
with open("input.txt") as infile, open("output.txt", "w") as outfile:
    outfile.write(infile.read())

数据库事务示例

import sqlite3
from contextlib import contextmanager


@contextmanager
def database_transaction(db_path: str):
    """数据库事务上下文管理器"""
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    try:
        yield cursor
        conn.commit()
        print("事务提交成功")
    except Exception as e:
        conn.rollback()
        print(f"事务回滚: {e}")
        raise
    finally:
        conn.close()


# 嵌套上下文管理器
def transfer_money(from_id: int, to_id: int, amount: float) -> None:
    with database_transaction("bank.db") as cursor:
        cursor.execute("UPDATE accounts SET balance = balance - ? WHERE id = ?",
                      (amount, from_id))
        cursor.execute("UPDATE accounts SET balance = balance + ? WHERE id = ?",
                      (amount, to_id))

自定义上下文管理器实战

实战 1:性能计时器

import time
from contextlib import contextmanager


@contextmanager
def timer(name: str = "block"):
    """代码块执行计时器"""
    start = time.perf_counter()
    print(f"[{name}] 开始执行...")
    try:
        yield
    finally:
        elapsed = time.perf_counter() - start
        print(f"[{name}] 耗时: {elapsed:.4f}s")


with timer("数据处理"):
    time.sleep(0.5)
    result = sum(i * i for i in range(10_000_000))
    print(f"计算结果: {result}")
# 输出:
# [数据处理] 开始执行...
# 计算结果: 333333283333335000000
# [数据处理] 耗时: 0.5234s

实战 2:临时目录切换

import os
from pathlib import Path
from contextlib import contextmanager


@contextmanager
def change_directory(target_dir: str):
    """临时切换工作目录"""
    original_dir = os.getcwd()
    try:
        os.chdir(target_dir)
        print(f"切换到目录: {target_dir}")
        yield
    finally:
        os.chdir(original_dir)
        print(f"恢复目录: {original_dir}")


# 使用示例
with change_directory(Path.home() / "Documents"):
    print(f"当前目录: {os.getcwd()}")
    # 在这里做文件操作...

print(f"当前目录: {os.getcwd()}")  # 已恢复

实战 3:线程锁上下文

import threading
from contextlib import contextmanager


class SharedCounter:
    """线程安全的计数器"""

    def __init__(self):
        self.value = 0
        self._lock = threading.Lock()

    def increment(self) -> None:
        with self._lock:  # 锁的上下文管理
            self.value += 1

    def decrement(self) -> None:
        with self._lock:
            self.value -= 1

    @contextmanager
    def batch_update(self):
        """批量更新上下文(一次加锁,多次操作)"""
        with self._lock:
            print("开始批量更新")
            yield
            print("批量更新结束")


counter = SharedCounter()

# 使用锁上下文
def worker():
    for _ in range(10000):
        counter.increment()

threads = [threading.Thread(target=worker) for _ in range(10)]
for t in threads:
    t.start()
for t in threads:
    t.join()

print(f"最终结果 (期望 100000): {counter.value}")

# 批量更新
with counter.batch_update():
    counter.value += 10
    counter.value += 20
    counter.value += 30

实战 4:临时环境变量

import os
from contextlib import contextmanager


@contextmanager
def set_env(**environ: str):
    """临时设置环境变量"""
    old_env = {}
    try:
        for key, value in environ.items():
            old_env[key] = os.environ.get(key)
            os.environ[key] = value
            print(f"设置 {key}={value}")
        yield
    finally:
        for key in environ:
            if old_env[key] is None:
                del os.environ[key]
            else:
                os.environ[key] = old_env[key]
            print(f"恢复 {key}={old_env[key]}")


with set_env(DATABASE_URL="sqlite:///test.db", DEBUG="true"):
    print(f"DATABASE_URL: {os.environ['DATABASE_URL']}")
    print(f"DEBUG: {os.environ['DEBUG']}")

print(f"DATABASE_URL: {os.environ.get('DATABASE_URL')}")  # 已恢复

实战 5:连接池管理

from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime


@dataclass
class Connection:
    """模拟数据库连接"""
    id: int
    created_at: datetime = datetime.now()

    def query(self, sql: str) -> str:
        return f"[连接 #{self.id}] 执行: {sql}"

    def close(self) -> None:
        print(f"关闭连接 #{self.id}")


class ConnectionPool:
    """数据库连接池"""

    def __init__(self, min_size: int = 2, max_size: int = 10):
        self.min_size = min_size
        self.max_size = max_size
        self._idle: list[Connection] = [Connection(i) for i in range(min_size)]
        self._active: set[Connection] = set()
        self._counter = min_size

    @contextmanager
    def get_connection(self) -> Connection:
        """获取连接(使用后自动归还)"""
        conn = self._acquire()
        try:
            yield conn
        finally:
            self._release(conn)

    def _acquire(self) -> Connection:
        if self._idle:
            conn = self._idle.pop()
            self._active.add(conn)
            print(f"从连接池获取连接 #{conn.id}")
            return conn
        if len(self._active) < self.max_size:
            self._counter += 1
            conn = Connection(self._counter)
            self._active.add(conn)
            print(f"创建新连接 #{conn.id}")
            return conn
        raise RuntimeError("连接池已满")

    def _release(self, conn: Connection) -> None:
        self._active.discard(conn)
        if len(self._idle) < self.max_size:
            self._idle.append(conn)
            print(f"归还连接 #{conn.id} 到连接池")
        else:
            conn.close()


pool = ConnectionPool(min_size=2)

# 使用连接
with pool.get_connection() as conn:
    result = conn.query("SELECT * FROM users")
    print(result)

with pool.get_connection() as conn:
    result = conn.query("SELECT * FROM orders")
    print(result)

类实现 vs 装饰器实现

两种实现上下文管理器的方式各有优劣:

对比项 类实现 @contextmanager
复杂度 稍复杂 简洁
异常处理 灵活(通过返回值控制抑制) 有限(只能通过 try/except)
需要记住状态 适合(self 可以存储状态) 不太适合
可读性 结构清晰 非常简洁
需要单独的类 否(函数即可)

选择原则

# 需要状态管理 → 类实现
class RetryContext:
    def __init__(self, max_retries: int = 3):
        self.max_retries = max_retries
        self.attempt = 0

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None and self.attempt < self.max_retries:
            self.attempt += 1
            print(f"重试 {self.attempt}/{self.max_retries}")
            return True  # 抑制异常,继续重试
        return False


# 简单场景 → @contextmanager
@contextmanager
def simple_timer():
    start = time.perf_counter()
    yield
    print(f"耗时: {time.perf_counter() - start:.4f}s")

__enter__ 返回值详解

__enter__ 的返回值通过 as 子句传递给变量:

class Connection:
    def __enter__(self):
        print("建立连接")
        return self  # 返回自身

    def __exit__(self, *args):
        print("关闭连接")

    def query(self, sql: str):
        return f"查询: {sql}"


class Config:
    def __enter__(self):
        print("加载配置")
        return {"debug": True, "port": 8080}  # 返回字典

    def __exit__(self, *args):
        print("保存配置")


# conn 是 Connection 实例
with Connection() as conn:
    print(conn.query("SELECT 1"))

# config 是字典
with Config() as config:
    print(config["debug"])

常见陷阱

陷阱 1:__exit__ 中返回了 True 但不应该

class SilentError:
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print(f"捕获到: {exc_val}")
        return True  # 抑制了异常!


with SilentError() as se:
    raise ValueError("严重错误!")
# 不会触发异常!这通常不是你想要的

陷阱 2:@contextmanager 生成器中没有 try/finally

@contextmanager
def bad_context():
    print("进入")
    yield "value"
    # 如果 yield 之前的代码或 with 块中抛异常,
    # yield 之后的代码不会执行!
    print("退出")  # 异常时不会执行!


@contextmanager
def good_context():
    print("进入")
    try:
        yield "value"
    finally:
        print("退出")  # 无论如何都会执行


with good_context() as v:
    print(v)
    raise ValueError("测试")

# 输出:
# 进入
# value
# 退出

陷阱 3:同时管理多个资源时忘记处理顺序

# 正确:先打开的资源后关闭
with open("a.txt") as a, open("b.txt") as b:
    pass  # a 先打开,b 后打开;退出时 b 先关闭,a 后关闭

# 等价于:
with open("a.txt") as a:
    with open("b.txt") as b:
        pass

小结

上下文管理器是 Python 中优雅管理资源的关键机制。本篇涵盖了:

  • with 语句:Python 推荐的安全资源管理方式
  • __enter__ / __exit__ 协议:类实现上下文管理器的基础
  • @contextmanager 装饰器:用生成器函数简化上下文管理器实现
  • contextlib 工具集closingsuppressnullcontextredirect_stdoutExitStack
  • 嵌套上下文管理器:同时管理多个资源
  • 实战示例:计时器、目录切换、线程锁、环境变量、连接池

何时使用上下文管理器

  • 文件操作
  • 数据库连接和事务
  • 线程锁
  • 网络连接
  • 临时环境修改
  • 计时和性能分析
  • 任何需要"进入-退出"模式的操作

掌握上下文管理器后,你就能写出更安全、更 Pythonic 的代码。

Summary: with 语句、enter/exit、@contextmanager、contextlib