Python中的上下文管理

上下文

定义

“上下文”（Context）指的是程序执行过程中的一种环境状态，它包含了当前执行代码的相关信息和资源。

可以理解为程序运行时的一个“场景”或“背景”，它决定了代码在特定环境下如何执行。

作用

资源管理：上下文可以管理程序运行时所需的资源，例如文件句柄、网络连接、数据库连接等。通过上下文管理器，可以在资源使用完毕后自动释放资源，避免资源泄漏。
状态维护：上下文可以保存程序的运行状态，例如线程的局部变量、函数调用栈信息等。这些状态信息可以帮助程序在复杂的执行环境中正确运行。
异常处理：上下文可以捕获和处理程序运行时的异常。通过上下文管理器，可以在异常发生时执行清理操作，确保程序的稳定性和可靠性。
环境隔离：上下文可以为程序提供一个独立的执行环境，避免不同部分的代码相互干扰。例如，在异步编程中，上下文变量可以确保协程之间的状态隔离。

‌上下文管理器

上下文管理器通过实现__enter__和__exit__方法管理资源，确保资源在使用前后自动完成初始化和清理操作。例如文件操作中，with open()自动关闭文件流‌。

__enter__() 方法当进入 with 语句块时，解释器会调用上下文管理器对象的 __enter__() 方法。这个方法定义了进入代码块时要执行的操作，并且可以返回一个对象。 __exit__() 方法当退出 with 语句块时，无论代码块是否出现异常，解释器都会调用上下文管理器对象的 __exit__() 方法。

它接收三个参数：

exc_type（异常类型）
exc_value（异常值）
traceback（异常的回溯信息），用于执行清理工作

with语句执行流程

使用 with 语句可以简化上下文管理器的使用，它会在执行 with 语句块之前自动调用 __enter__() 方法，并在执行结束后自动调用 __exit__() 方法

with 上下文管理器 as 变量:
    # 执行代码块

内置上下文管理器

文件操作

文件对象是一个内置的上下文管理器。使用 with 语句可以自动管理文件的打开和关闭，即使在读取或写入文件时发生异常也能确保文件被正确关闭。

with open('example.txt', 'r') as file:
    content = file.read()

线程锁

在多线程编程中，threading.Lock 也是一个上下文管理器，可以自动获取和释放锁，避免死锁的发生。

with threading.Lock():
    # 临界区操作

自定义上下文管理器

通过类实现

定义一个类并实现 __enter__() 和 __exit__() 方法。

 
class MyContextManager:
    def __enter__(self):
        print("Entering the context")
        return self
 
    def __exit__(self, exc_type, exc_value, traceback):
        print("Exiting the context")
        if exc_type:
            print(f"Exception occurred: {exc_type}, {exc_value}")
        return True  # 返回 True 表示异常被处理
 
with MyContextManager() as manager:
    print("Inside the context")

使用 contextlib 模块

contextlib 模块提供了 contextmanager 装饰器，可以更方便地创建上下文管理器。

from contextlib import contextmanager
 
@contextmanager
def file_opener(path):
    f = open(path, 'r')
    try:
        yield f
    finally:
        f.close()
 
with file_opener('data.txt') as f:
    content = f.read()‌:ml-citation{ref="2,8" data="citationList"}

高级特性

异常处理机制‌

__exit__方法返回True时，异常被静默处理；返回False或None时异常向外传播‌。

 
class SafeContext:
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type:
            print(f"异常捕获：{exc_val}")
            return True  # 阻止异常传播

嵌套上下文管理‌

支持多层with语句嵌套，例如同时管理文件和数据库连接‌。

with open('a.txt') as f1, open('b.txt') as f2:
    data = f1.read() + f2.read()

异步上下文

在 Python 3.5 中，随着 async 和 await 语法的引入，异步编程得到了更广泛的支持。

异步上下文管理器是这一背景下引入的一个特性，它允许在异步代码中使用上下文管理器，从而更好地管理异步资源的生命周期

__aenter__()：在进入 async with 语句块时被调用，用于初始化资源。
__aexit__()：在退出 async with 语句块时被调用，用于清理资源。

 
import asyncio
 
class AsyncContextManager:
    async def __aenter__(self):
        print("Entering the context")
        await asyncio.sleep(1)  # 模拟异步操作
        return self
 
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        print("Exiting the context")
        await asyncio.sleep(1)  # 模拟异步操作
 
async def main():
    async with AsyncContextManager() as manager:
        print("Inside the context")
 
asyncio.run(main())

异步文件操作：可以使用异步上下文管理器来管理异步文件操作。
异步数据库连接：在异步数据库操作中，可以使用异步上下文管理器来管理数据库连接。
异步锁和信号量：在异步并发编程中，可以使用异步上下文管理器来管理锁和信号量。给出代码和注释说明。

异步文件操作

在 Python 中，可以使用 aiofiles 库来实现异步文件操作。aiofiles 是一个支持异步上下文管理器的文件操作库。

import asyncio
import aiofiles  # 需要安装 aiofiles 库
 
async def read_file_async(file_path):
    # 使用异步上下文管理器打开文件
    async with aiofiles.open(file_path, mode='r') as file:
        content = await file.read()  # 异步读取文件内容
        print(content)
 
async def write_file_async(file_path, content):
    # 使用异步上下文管理器打开文件并写入内容
    async with aiofiles.open(file_path, mode='w') as file:
        await file.write(content)  # 异步写入文件内容
 
async def main():
    # 写入文件
    await write_file_async('example.txt', 'Hello, async world!')
    # 读取文件
    await read_file_async('example.txt')
 
# 运行异步主函数
asyncio.run(main())

aiofiles.open() 是一个异步上下文管理器，支持 async with 语法。在 __aenter__() 方法中，文件被异步打开。在 __aexit__() 方法中，文件被异步关闭。使用 await 关键字可以等待异步操作完成，避免阻塞事件循环。

异步数据库连接

使用 aiomysql 的代码示例，展示如何通过异步上下文管理器管理异步数据库连接。

import asyncio
import aiomysql
 
async def main():
    # 创建数据库连接
    # 使用 async with 语法创建连接，确保连接在退出时自动关闭
    async with aiomysql.connect(
        host='127.0.0.1',  # 数据库地址
        port=3306,         # 数据库端口
        user='your_username',  # 数据库用户名
        password='your_password',  # 数据库密码
        db='your_database',  # 数据库名称
        charset='utf8mb4'    # 字符集
    ) as conn:
        # 创建游标对象
        async with conn.cursor() as cursor:
            # 执行 SQL 查询
            await cursor.execute("SELECT * FROM your_table")
            # 获取查询结果
            result = await cursor.fetchall()
            for row in result:
                print(row)
 
# 运行异步主函数
asyncio.run(main())

通过使用 async with 语法，可以显著简化数据库连接和游标的管理。具体优势如下：

自动资源管理：连接和游标在退出上下文时会自动关闭，无需手动管理，减少了资源泄漏的风险。
代码清晰：async with 语法使得代码更加简洁，逻辑更加清晰。
异常安全：即使在执行 SQL 查询或获取结果时发生异常，连接和游标也会被正确关闭。

异步锁和信号量

在异步并发编程中，可以使用 asyncio.Lock 或 asyncio.Semaphore 来管理并发访问。这些对象都支持异步上下文管理器。

import asyncio
 
async def critical_section(lock):
    async with lock:  # 使用异步上下文管理器获取锁
        print(f"{asyncio.current_task().get_name()} entered the critical section")
        await asyncio.sleep(1)  # 模拟耗时操作
        print(f"{asyncio.current_task().get_name()} exiting the critical section")
 
async def main():
    lock = asyncio.Lock()  # 创建一个异步锁
    tasks = [
        asyncio.create_task(critical_section(lock), name="Task 1"),
        asyncio.create_task(critical_section(lock), name="Task 2"),
    ]
    await asyncio.gather(*tasks)  # 等待所有任务完成
 
# 运行异步主函数
asyncio.run(main())

asyncio.Lock 是一个异步锁，支持异步上下文管理器。在 __aenter__() 方法中，锁被异步获取。在 __aexit__() 方法中，锁被异步释放。使用 async with lock 可以确保锁在进入和退出上下文时被正确管理。

资源释放的确定性

即使代码块中发生return或sys.exit()，__exit__仍会执行，确保资源释放‌。

contextvars

说明

ContextVar用于声明与当前执行上下文（如线程、协程）绑定的变量，确保不同并发任务的状态隔离。类似于 Java中的ThreadLocal。

上下文变量是一种特殊的变量，它的值会随着执行上下文（如协程、线程）的变化而变化，这使得在异步编程和多线程环境中能够更方便地管理和传递上下文信息。

定义

使用 contextvars.ContextVar 创建上下文变量。每个上下文变量都有一个唯一的名称，并且可以有一个默认值。

import contextvars
request_id = contextvars.ContextVar('name', default=None)

name参数‌：用于调试和日志标识变量。
‌default参数‌：当上下文中未设置值时，var.get()返回默认值‌

基础操作

置值‌：token = var.set(100)，返回Token对象用于恢复之前状态‌。
取值‌：current_value = var.get()，若未设置则返回默认值或抛出LookupError。
‌重置值‌：var.reset(token)，将变量恢复至Token对应的状态‌。

应用场景

异步编程：在异步任务中使用上下文变量可以确保在一个任务中设置的变量值不会影响其他任务。

import contextvars
import asyncio
 
var = contextvars.ContextVar('var', default='default_value')
 
async def worker(worker_id):
    print(f'Worker {worker_id} 设置前的值: {var.get()}')
    var.set(f'value_for_worker_{worker_id}')
    print(f'Worker {worker_id} 设置后的值: {var.get()}')
    await asyncio.sleep(1)
    print(f'Worker {worker_id} 结束时的值: {var.get()}')
 
async def main():
    async with asyncio.TaskGroup() as tg:
        tg.create_task(worker(1))
        tg.create_task(worker(2))
 
asyncio.run(main())

日志追踪：在分布式系统中，每个请求可能会经过多个服务和组件。使用上下文变量可以在整个请求处理过程中传递唯一的追踪 ID，方便日志记录和问题排查。 ‌Web请求上下文管理‌，在Web框架（如Flask）中存储当前请求的认证信息或用户ID，确保不同请求状态隔离。‌在Fastapi框架中也可以用来实现传递请求头中的参数，比如说多租户场景下的tenant_id。

 
import logging
import contextvars
 
trace_id_var = contextvars.ContextVar('trace_id', default='unknown')
 
class TraceIdFilter(logging.Filter):
    def filter(self, record):
        record.trace_id = trace_id_var.get()
        return True
 
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
logger.addFilter(TraceIdFilter())
 
def process_request():
    token = trace_id_var.set('123456')
    try:
        logger.info("处理请求")
    finally:
        trace_id_var.reset(token)
 
process_request()

工具类封装

 
 
from contextvars import ContextVar, Token
from typing import (
    Any,
    Generic,
    TypeVar,
    Optional,
    Dict,
    Callable,
    Union,
    overload,
)
from functools import wraps
import uuid
 
T = TypeVar("T")
KT = TypeVar("KT")
 
class ContextVarManager(Generic[T]):
    """
    线程/协程安全的上下文变量管理器
    特性：
    1. 强类型校验 (通过mypy检查)
    2. 支持动态默认值
    3. 嵌套上下文恢复
    4. 唯一标识符生成
    5. 自动注册管理
    """
 
    _registry: Dict[str, "ContextVarManager"] = {}
 
    def __init__(
        self,
        name: Optional[str] = None,
        *,
        default: Union[T, Callable[[], T]] = None,
        enforce_type: bool = True,
    ):
        self._name = name or f"ctx_var_{uuid.uuid4().hex}"
        self._var: ContextVar[T] = ContextVar(self._name)
        self._default = default
        self._enforce_type = enforce_type
        self.__class__._registry[self._name] = self
 
    @overload
    def get(self) -> T:
        ...
 
    @overload
    def get(self, default: KT) -> Union[T, KT]:
        ...
 
    def get(self, default=...):
        """获取当前上下文值，支持动态默认值"""
        try:
            return self._var.get()
        except LookupError:
            if default is not ...:
                return default
            if self._default is None:
                raise
            return self._default() if callable(self._default) else self._default
 
    def set(self, value: T) -> Token:
        """设置上下文值，可选类型校验"""
        if self._enforce_type and not isinstance(value, self.__orig_class__.__args__):
            expected_type = self.__orig_class__.__args__
            raise TypeError(
                f"Expected type {expected_type}, got {type(value).__name__}"
            )
        return self._var.set(value)
 
    def reset(self, token: Token) -> None:
        """恢复到指定token状态"""
        self._var.reset(token)
 
    @classmethod
    def get_all(cls) -> Dict[str, Any]:
        """获取所有注册的上下文变量当前值"""
        return {
            name: manager.get(None)
            for name, manager in cls._registry.items()
        }
 
    def __enter__(self) -> "ContextVarManager":
        """支持with语法快速设置临时值"""
        self._prev_token = self._var.set(self._default)
        return self
 
    def __exit__(self, *args: Any) -> None:
        self._var.reset(self._prev_token)
 
    @classmethod
    def context_wrapper(cls, **context_values: Any):
        """
        装饰器：为函数调用创建独立上下文
        示例：
        @ContextVarManager.context_wrapper(user_id=123)
        def my_func():
            print(current_user.get())  # 123
        """
        def decorator(func: Callable):
            @wraps(func)
            def wrapper(*args, **kwargs):
                tokens = []
                try:
                    for name, value in context_values.items():
                        manager = cls._registry[name]
                        tokens.append(manager.set(value))
                    return func(*args, **kwargs)
                finally:
                    for manager, token in zip(cls._registry.values(), reversed(tokens)):
                        manager.reset(token)
            return wrapper
        return decorator
 
# 类型安全的变量声明示例
current_user: ContextVarManager[str] = ContextVarManager(default="guest")
request_id: ContextVarManager[str] = ContextVarManager(default_factory=lambda: uuid.uuid4().hex)
debug_mode: ContextVarManager[bool] = ContextVarManager(default=False)

主流框架中的上下文

Flask

Flask 通过 ‌请求上下文‌ 和 ‌应用上下文‌ 实现线程/协程隔离的全局状态管理

请求上下文

‌应用场景‌：存储 HTTP 请求的头部、参数、Cookie 等数据，支持多请求并发处理‌。

# Flask 源码片段（flask/ctx.py）  
class RequestContext:  
    def __init__(self, app, environ, request=None):  
        self.app = app  
        self.request = request or Request(environ)  
        self.session = None  # 会话数据初始化  
 
    def push(self):  
        _request_ctx_stack.push(self)  # 压入请求上下文栈  
 
# 使用 LocalProxy 代理访问当前请求  
request = LocalProxy(partial(_lookup_req_object, "request"))  
session = LocalProxy(partial(_lookup_req_object, "session"))

每个请求通过 push() 方法将上下文压入线程隔离的栈 (_request_ctx_stack)‌。
LocalProxy 动态代理确保不同请求访问独立的 request 和 session。

‌应用上下文

存储应用配置、数据库连接等全局资源‌。

# Flask 应用上下文管理（flask/ctx.py）  
class AppContext:  
    def __init__(self, app):  
        self.app = app  
        self.g = app.app_ctx_globals_class()  
 
    def push(self):  
        _app_ctx_stack.push(self)  # 压入应用上下文栈  
 
current_app = LocalProxy(_find_app)  
g = LocalProxy(partial(_lookup_app_object, "g"))

g 对象用于请求周期内的临时存储（如数据库连接）‌。
current_app 指向当前激活的 Flask 应用实例，支持多应用共存‌。

FastAPI

FastAPI 结合 ‌依赖注入‌ 和 contextvars 实现异步友好的上下文管理。

在中间件或路由中共享数据库连接或用户认证状态‌。

from fastapi import Depends, Request  
from contextvars import ContextVar  
 
db_connection: ContextVar[Connection] = ContextVar("db_conn")  
 
async def get_db(request: Request):  
    conn = acquire_connection()  
    token = db_connection.set(conn)  # 设置上下文变量  
    try:  
        yield conn  
    finally:  
        db_connection.reset(token)  
        release_connection(conn)  
 
@app.get("/items")  
async def read_items(db: Connection = Depends(get_db)):  
    return {"data": db.query(Item)}

使用 ContextVar 确保异步任务中数据库连接的隔离性‌。
依赖注入通过 yield 实现资源的自动释放（类似上下文管理器）‌。

Celery

应用场景‌：分布式任务中跟踪请求链路的唯一标识‌。

from celery import Celery, Task  
from contextvars import ContextVar  
 
request_id: ContextVar[str] = ContextVar("request_id")  
 
class ContextTask(Task):  
    def __call__(self, *args, **kwargs):  
        token = request_id.set(kwargs.pop("request_id"))  
        try:  
            return super().__call__(*args, **kwargs)  
        finally:  
            request_id.reset(token)  
 
app = Celery()  
app.Task = ContextTask  
 
@app.task  
def process_data(data):  
    print(f"Processing {data} with ID: {request_id.get()}")

自定义 Task 类在任务执行前注入上下文变量（如 request_id）‌。结合 ContextVar 实现多任务并发下的状态隔离‌。

Django

请求处理中间件‌

‌应用场景‌：在请求周期内添加用户身份或时区信息‌。

# middleware.py  
class UserMiddleware:  
    def __init__(self, get_response):  
        self.get_response = get_response  
 
    def __call__(self, request):  
        request.user = get_user(request)  # 设置请求上下文  
        response = self.get_response(request)  
        return response  
 
# settings.py  
MIDDLEWARE = [  
    'middleware.UserMiddleware',  
]

中间件在请求处理前设置 request.user，全局可用‌。

模板上下文处理器‌

‌应用场景‌：向模板中注入全局变量（如站点配置）‌。

# context_processors.py  
def site_info(request):  
    return {  
        "SITE_NAME": "MySite",  
        "CURRENT_YEAR": datetime.now().year  
    }  
 
# settings.py  
TEMPLATES = [{  
    'OPTIONS': {  
        'context_processors': [  
            'context_processors.site_info',  
        ],  
    },  
}]

上下文处理器自动向所有模板传递预定义变量‌。