Python 上下文管理器完全指南：从 with 语句到高级实战

with open("data.txt", "r") as f:
    content = f.read()

class MyContext:
    def __enter__(self):
        print("进入上下文")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        print("退出上下文")
        # 返回 False 表示不处理异常，让异常继续传播
        # 返回 True 表示异常已被处理，不再传播
        return False

with MyContext() as ctx:
    print("正在执行")
    # 即使这里抛出异常，__exit__ 也会被调用

class DatabaseConnection:
    def __init__(self, conn_string: str):
        self.conn_string = conn_string
        self.connection = None
    
    def __enter__(self):
        print(f"连接数据库: {self.conn_string}")
        self.connection = self._connect()
        return self.connection
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type:
            print(f"发生异常: {exc_val}")
            self.connection.rollback()
        else:
            self.connection.commit()
        self.connection.close()
        print("数据库连接已关闭")
        return False  # 不吞掉异常
    
    def _connect(self):
        # 模拟数据库连接
        return {"conn": self.conn_string, "status": "connected"}

with DatabaseConnection("postgresql://localhost/mydb") as conn:
    # 执行数据库操作
    # 如果抛出异常，自动回滚
    # 如果正常结束，自动提交
    pass  # 无论怎样，连接都会关闭

from contextlib import contextmanager

@contextmanager
def database_connection(conn_string: str):
    print(f"连接数据库: {conn_string}")
    conn = {"conn": conn_string, "status": "connected"}
    try:
        yield conn  # __enter__ 的返回值
    except Exception as e:
        print(f"回滚事务: {e}")
        conn["rollback"] = True
        raise  # 重新抛出异常
    else:
        conn["commit"] = True
        print("提交事务")
    finally:
        conn["status"] = "closed"
        print("数据库连接已关闭")

with database_connection("postgresql://localhost/mydb") as conn:
    # 业务逻辑
    pass

from contextlib import suppress
import os

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

# 上下文管理器写法
with suppress(FileNotFoundError):
    os.remove("temp.txt")

with suppress(FileNotFoundError, PermissionError):
    os.remove("/etc/protected/config")

from contextlib import redirect_stdout
import io

buf = io.StringIO()
with redirect_stdout(buf):
    print("这条信息不会打印到终端")
    print("而是被收集到 buf 中")

output = buf.getvalue()
print(f"捕获的内容: {output}")

def test_verbose_function():
    buf = io.StringIO()
    with redirect_stdout(buf):
        my_function(verbose=True)
    assert "已完成" in buf.getvalue()

from contextlib import chdir
import os

cwd = os.getcwd()
print(f"当前目录: {cwd}")

with chdir("/tmp"):
    print(f"临时在: {os.getcwd()}")
    # 做一些文件操作

print(f"回到: {os.getcwd()}")

old_cwd = os.getcwd()
try:
    os.chdir("/tmp")
    # 业务逻辑
finally:
    os.chdir(old_cwd)

from contextlib import closing
from urllib.request import urlopen

with closing(urlopen("https://python.org")) as page:
    for line in page:
        if b"Python" in line:
            print(line)

from decimal import Decimal, localcontext

d = Decimal("1") / Decimal("7")
print(d)  # 默认精度: 0.1428571428571428571428571429

with localcontext() as ctx:
    ctx.prec = 3
    d2 = Decimal("1") / Decimal("7")
    print(d2)  # 0.143

print(Decimal("1") / Decimal("7"))  # 恢复默认精度

import threading

lock = threading.Lock()
shared_resource = []

def worker():
    with lock:
        # 只有拿到锁的线程能进入这个代码块
        shared_resource.append(threading.current_thread().name)

from contextlib import contextmanager

@contextmanager
def transaction(connection):
    """数据库事务上下文管理器"""
    try:
        print("开始事务")
        yield connection
        print("提交事务")
        connection.commit()
    except Exception:
        print("回滚事务")
        connection.rollback()
        raise

@contextmanager
def db_cursor(connection):
    """数据库游标上下文管理器"""
    cursor = connection.cursor()
    try:
        yield cursor
    finally:
        cursor.close()

# 嵌套使用
with DatabaseConnection("postgresql://localhost/mydb") as conn:
    with transaction(conn):
        with db_cursor(conn) as cursor:
            cursor.execute("UPDATE users SET name = ? WHERE id = ?", ("Alice", 1))
            cursor.execute("UPDATE users SET name = ? WHERE id = ?", ("Bob", 2))

# Python 3.10+ 可以用括号分组
with (
    DatabaseConnection("postgresql://localhost/mydb") as conn,
    db_cursor(conn) as cursor,
    transaction(conn),
):
    cursor.execute("UPDATE users SET name = ? WHERE id = ?", ("Alice", 1))

from contextlib import ExitStack

def process_files(file_names: list[str]):
    """同时打开多个文件"""
    with ExitStack() as stack:
        files = [
            stack.enter_context(open(fname, "r"))
            for fname in file_names
        ]
        # 所有文件在退出 with 块时自动关闭
        return [f.read() for f in files]

from contextlib import ExitStack, contextmanager

@contextmanager
def measure_time(label: str):
    import time
    start = time.time()
    try:
        yield
    finally:
        elapsed = time.time() - start
        print(f"{label}: {elapsed:.3f}s")

def complex_operation(debug: bool = False):
    with ExitStack() as stack:
        if debug:
            stack.enter_context(measure_time("复杂操作"))
            stack.enter_context(redirect_stdout(open("debug.log", "w")))
        
        # 核心逻辑
        result = do_heavy_computation()
        return result

def deploy_app():
    """部署应用的流程，任何一步失败都要回滚"""
    with ExitStack() as stack:
        # 注册回滚回调
        def rollback_nginx():
            print("回滚: 恢复 nginx 配置")
        
        def rollback_db():
            print("回滚: 恢复数据库版本")
        
        # 每个操作成功后，注册对应的回滚函数
        print("步骤1: 更新数据库")
        migrate_database()
        stack.callback(rollback_db)
        
        print("步骤2: 部署新代码")
        deploy_code()
        stack.callback(lambda: print("回滚: 恢复代码版本"))
        
        print("步骤3: 更新 nginx 配置")
        update_nginx()
        stack.callback(rollback_nginx)
        
        # 如果所有步骤成功，弹出所有回滚（不执行）
        stack.pop_all()
        print("部署成功!")
    
    # 如果任何一步抛出异常，ExitStack 会按注册的逆序执行回调

import time
from contextlib import contextmanager
from collections.abc import Generator

@contextmanager
def timer(label: str = "耗时") -> Generator[None, None, None]:
    start = time.perf_counter()
    try:
        yield
    finally:
        elapsed = time.perf_counter() - start
        print(f"{label}: {elapsed:.4f}秒")

# 使用
with timer("数据处理"):
    result = [i ** 2 for i in range(10_000_000)]

import contextvars

_depth = contextvars.ContextVar("timer_depth", default=0)

@contextmanager
def timer(label: str = "耗时"):
    depth = _depth.get()
    _depth.set(depth + 1)
    prefix = "  " * depth
    start = time.perf_counter()
    try:
        yield
    finally:
        _depth.set(depth)
        elapsed = time.perf_counter() - start
        print(f"{prefix}{label}: {elapsed:.4f}秒")

# 嵌套使用
with timer("外部"):
    with timer("内部A"):
        time.sleep(0.1)
    with timer("内部B"):
        time.sleep(0.2)
# 输出：
#   内部A: 0.1001秒
#   内部B: 0.2001秒
# 外部: 0.3010秒

import queue
from contextlib import contextmanager

class ConnectionPool:
    def __init__(self, size: int = 5):
        self._pool = queue.Queue(maxsize=size)
        for i in range(size):
            self._pool.put(self._create_connection(i))
    
    def _create_connection(self, idx: int):
        return {"id": idx, "active": False}
    
    @contextmanager
    def get_connection(self):
        conn = self._pool.get()
        conn["active"] = True
        try:
            yield conn
        finally:
            conn["active"] = False
            self._pool.put(conn)

pool = ConnectionPool(3)

with pool.get_connection() as conn:
    print(f"使用连接: {conn['id']}")
# 连接自动归还到池中

import os
from contextlib import contextmanager

@contextmanager
def set_env(**environ: str):
    """临时设置环境变量"""
    old_environ = dict(os.environ)
    os.environ.update(environ)
    try:
        yield
    finally:
        os.environ.clear()
        os.environ.update(old_environ)

# 使用
with set_env(DATABASE_URL="sqlite:///test.db", DEBUG="1"):
    assert os.environ["DATABASE_URL"] == "sqlite:///test.db"
    assert os.environ["DEBUG"] == "1"
# 退出后自动恢复

from dataclasses import dataclass, field
from contextlib import contextmanager

@dataclass
class AppConfig:
    debug: bool = False
    db_url: str = "postgresql://localhost/prod"
    
    @contextmanager
    def override(self, **kwargs):
        """临时覆盖配置"""
        old_values = {}
        for key, value in kwargs.items():
            if hasattr(self, key):
                old_values[key] = getattr(self, key)
                setattr(self, key, value)
        try:
            yield self
        finally:
            for key, value in old_values.items():
                setattr(self, key, value)

config = AppConfig()
print(config.db_url)  # postgresql://localhost/prod

with config.override(db_url="sqlite:///test.db", debug=True):
    print(config.db_url)  # sqlite:///test.db
    print(config.debug)   # True

print(config.db_url)  # postgresql://localhost/prod（恢复）

# 错误理解
with open("nonexistent.txt", "r") as f:
    content = f.read()
# FileNotFoundError 会直接传播出去

# 正确理解：with 只能保证资源释放，不能保证异常不传播

@contextmanager
def bad_context():
    yield "resource"
    # 如果 yield 之后的代码块抛了异常，这里永远执行不到！
    print("清理资源")  # ❌ 不会执行

# 正确写法
@contextmanager
def good_context():
    print("获取资源")
    try:
        yield "resource"
    finally:
        print("清理资源")  # ✅ 无论是否异常都会执行

class SilentContext:
    def __exit__(self, exc_type, exc_val, exc_tb):
        return True  # 🚨 返回 True 会吞掉异常！

with SilentContext():
    raise ValueError("程序崩溃了！")
# 这个异常被静默吃掉了——你不会知道它发生了

# ❌ 错误：普通 contextmanager 不能用在异步函数中
@contextmanager
async def async_context():  # 这样不行！
    yield

# ✅ 正确：使用 asynccontextmanager
from contextlib import asynccontextmanager

@asynccontextmanager
async def async_context():
    print("异步获取资源")
    try:
        yield "资源"
    finally:
        print("异步释放资源")

# 使用
async def main():
    async with async_context() as resource:
        print(f"使用: {resource}")

import time
import logging
from contextlib import contextmanager
from typing import Type, Callable, Optional

logger = logging.getLogger(__name__)

class RetryExhaustedError(Exception):
    """所有重试均失败的异常"""
    pass

@contextmanager
def retry_context(
    max_attempts: int = 3,
    delay: float = 1.0,
    backoff: float = 2.0,
    exceptions: tuple[Type[Exception], ...] = (Exception,),
    on_retry: Optional[Callable[[int, Exception], None]] = None,
):
    """
    可重试的上下文管理器。
    
    用法:
        with retry_context(max_attempts=3) as attempt:
            with attempt():
                # 可能会失败的操作
                fetch_data_from_api()
    """
    
    class RetryAttempt:
        def __enter__(self):
            return self
        
        def __exit__(self, exc_type, exc_val, exc_tb):
            if exc_type is None or not issubclass(exc_type, exceptions):
                # 没有异常，或者是不需要重试的异常类型
                return False
            
            nonlocal max_attempts, delay
            
            attempt_number = max_attempts  # 从当前剩余次数倒推
            max_attempts -= 1
            
            if max_attempts > 0:
                if on_retry:
                    on_retry(attempt_number, exc_val)
                
                logger.warning(
                    "操作失败（剩余 %d 次重试）: %s",
                    max_attempts, exc_val
                )
                time.sleep(delay)
                delay *= backoff  # 指数退避
                return True  # 抑制异常，让循环继续
            else:
                raise RetryExhaustedError(
                    f"重试 {attempt_number} 次后仍然失败"
                ) from exc_val
    
    attempts = [RetryAttempt() for _ in range(max_attempts)]
    
    class RetryController:
        def __init__(self):
            self._attempts = iter(attempts)
        
        def __call__(self) -> RetryAttempt:
            return next(self._attempts)
    
    yield RetryController()


# ===== 使用示例 =====

import random

def unstable_api_call():
    """模拟一个偶尔失败的外部 API 调用"""
    if random.random() < 0.7:  # 70% 概率失败
        raise ConnectionError("网络超时")
    return {"status": "ok", "data": [1, 2, 3]}

# 使用重试上下文
with retry_context(
    max_attempts=5,
    delay=0.5,
    backoff=1.5,
    on_retry=lambda attempt, e: print(f"第 {attempt} 次重试，原因: {e}"),
) as attempt:
    for _ in range(5):  # 最多 5 次
        with attempt():
            result = unstable_api_call()
            print(f"成功: {result}")
            break

class AsyncDatabase:
    async def __aenter__(self):
        print("异步连接数据库")
        self.conn = await create_async_connection()
        return self.conn
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.conn.close()
        print("异步关闭连接")

# 使用
async def main():
    async with AsyncDatabase() as db:
        await db.query("SELECT * FROM users")

from contextlib import asynccontextmanager
import aiohttp

@asynccontextmanager
async def http_session(base_url: str):
    session = aiohttp.ClientSession(base_url)
    try:
        yield session
    finally:
        await session.close()

async def main():
    async with http_session("https://api.example.com") as session:
        async with session.get("/users") as resp:
            data = await resp.json()
            print(data)

from contextlib import AsyncExitStack

async def complex_async_workflow():
    async with AsyncExitStack() as stack:
        # 动态注册异步资源
        db = await stack.enter_async_context(
            AsyncDatabase()
        )
        session = await stack.enter_async_context(
            http_session("https://api.example.com")
        )
        # 也支持同步资源
        f = stack.enter_context(open("log.txt", "w"))
        
        # 所有资源会在退出 async with 时
        # 按注册逆序自动清理
        return await process(db, session, f)