Python 上下文管理器进阶：自定义实现与性能优化

Python 上下文管理器进阶：自定义实现与性能优化

1. 技术分析

1.1 上下文管理器定义

上下文管理器是实现了__enter__和__exit__方法的对象，用于管理资源的获取和释放：

with context_manager as resource: # 使用资源 pass # 资源自动释放

1.2 上下文管理器协议

1.3 上下文管理器分类

2. 核心功能实现

2.1 基础上下文管理器

import os class FileManager: def __init__(self, filename, mode='r'): self.filename = filename self.mode = mode self.file = None def __enter__(self): self.file = open(self.filename, self.mode) return self.file def __exit__(self, exc_type, exc_val, exc_tb): if self.file: self.file.close() return False class DatabaseConnection: def __init__(self, db_url): self.db_url = db_url self.connection = None def __enter__(self): self.connection = self._connect() return self.connection def __exit__(self, exc_type, exc_val, exc_tb): if self.connection: if exc_type is None: self.connection.commit() else: self.connection.rollback() self.connection.close() return False def _connect(self): import sqlite3 return sqlite3.connect(self.db_url) class Timer: def __init__(self, label=''): self.label = label self.start_time = None self.duration = None def __enter__(self): self.start_time = time.time() return self def __exit__(self, exc_type, exc_val, exc_tb): self.duration = time.time() - self.start_time print(f"{self.label} took {self.duration:.4f} seconds") return False

2.2 使用 @contextmanager 装饰器

from contextlib import contextmanager import threading @contextmanager def lock_resource(lock): lock.acquire() try: yield finally: lock.release() @contextmanager def temporary_change(obj, attr, value): original = getattr(obj, attr) setattr(obj, attr, value) try: yield finally: setattr(obj, attr, original) @contextmanager def suppress(*exceptions): try: yield except exceptions: pass @contextmanager def nullcontext(enter_result=None): yield enter_result class ThreadSafeResource: def __init__(self): self._lock = threading.Lock() self._data = {} @contextmanager def access(self, key): with self._lock: if key not in self._data: self._data[key] = [] yield self._data[key]

2.3 嵌套上下文管理器

from contextlib import ExitStack class MultiResourceManager: def __init__(self): self._resources = [] def add_resource(self, resource): self._resources.append(resource) def __enter__(self): return [r.__enter__() for r in self._resources] def __exit__(self, exc_type, exc_val, exc_tb): for resource in reversed(self._resources): try: resource.__exit__(exc_type, exc_val, exc_tb) except Exception: pass return False class TransactionManager: def __init__(self, db): self.db = db self._savepoints = [] def __enter__(self): self.db.begin() return self def __exit__(self, exc_type, exc_val, exc_tb): if exc_type is None: self.db.commit() else: self.db.rollback() def savepoint(self, name): self._savepoints.append(name) self.db.execute(f"SAVEPOINT {name}") def rollback_to_savepoint(self, name): self.db.execute(f"ROLLBACK TO SAVEPOINT {name}") def nested_transactions(): with ExitStack() as stack: conn1 = stack.enter_context(DatabaseConnection('db1.db')) conn2 = stack.enter_context(DatabaseConnection('db2.db')) try: conn1.execute('INSERT INTO table1 VALUES (1)') conn2.execute('INSERT INTO table2 VALUES (2)') except Exception: raise

2.4 异步上下文管理器

import asyncio from contextlib import asynccontextmanager class AsyncFileManager: def __init__(self, filename, mode='r'): self.filename = filename self.mode = mode self.file = None async def __aenter__(self): self.file = await asyncio.to_thread(open, self.filename, self.mode) return self.file async def __aexit__(self, exc_type, exc_val, exc_tb): if self.file: await asyncio.to_thread(self.file.close) return False @asynccontextmanager async def async_lock(lock): await lock.acquire() try: yield finally: lock.release() @asynccontextmanager async def async_timer(label=''): start = time.time() try: yield finally: elapsed = time.time() - start print(f"{label} took {elapsed:.4f}s") class AsyncDatabasePool: def __init__(self, max_connections=10): self._pool = [] self._max_connections = max_connections self._lock = asyncio.Lock() async def get_connection(self): async with self._lock: if self._pool: return self._pool.pop() if len(self._pool) < self._max_connections: return await self._create_connection() raise RuntimeError("Pool exhausted") async def release_connection(self, conn): async with self._lock: if len(self._pool) < self._max_connections: self._pool.append(conn) @asynccontextmanager async def connection(self): conn = await self.get_connection() try: yield conn finally: await self.release_connection(conn)

3. 性能对比

3.1 上下文管理器类型性能对比

3.2 资源管理开销对比

3.3 嵌套上下文性能

4. 最佳实践

4.1 上下文管理器设计模式

class ResourcePool: def __init__(self, create_func, max_size=10): self._create_func = create_func self._max_size = max_size self._available = [] self._in_use = set() self._lock = threading.Lock() def acquire(self): with self._lock: if self._available: resource = self._available.pop() self._in_use.add(id(resource)) return resource if len(self._in_use) < self._max_size: resource = self._create_func() self._in_use.add(id(resource)) return resource raise RuntimeError("Pool exhausted") def release(self, resource): with self._lock: if id(resource) in self._in_use: self._in_use.remove(id(resource)) if len(self._available) < self._max_size: self._available.append(resource) @contextmanager def get(self): resource = self.acquire() try: yield resource finally: self.release(resource)

4.2 上下文管理器与异常处理

class RobustContextManager: def __init__(self): self._resources = [] def register(self, resource, cleanup_func): self._resources.append((resource, cleanup_func)) def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): errors = [] for resource, cleanup in reversed(self._resources): try: cleanup(resource) except Exception as e: errors.append(e) if errors and not exc_type: raise ExceptionGroup("Multiple cleanup errors", errors) return False

5. 总结

上下文管理器是 Python 资源管理的核心机制：

1. 类实现提供最完整的控制能力
2. @contextmanager简化简单场景的实现
3. 异步上下文管理器支持 async/await 模式
4. ExitStack优雅处理动态数量的资源

对比数据如下：

• 上下文管理器引入的开销约为 0.03μs
• 异步上下文管理器开销约为同步版本的 5-10 倍
• 嵌套上下文的性能下降呈线性增长
• 在资源密集型场景中，上下文管理器的开销可忽略不计