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 资源管理的核心机制：

对比数据如下：

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