本文目录导读:
我来展示Python上下文管理器的几种实现方式,包含经典案例和实际应用场景。
使用类实现(__enter__ 和 __exit__)
文件操作案例
class FileManager:
def __init__(self, filename, mode='r', encoding='utf-8'):
self.filename = filename
self.mode = mode
self.encoding = encoding
self.file = None
def __enter__(self):
# 打开文件并返回文件对象
self.file = open(self.filename, self.mode, encoding=self.encoding)
print(f"打开文件: {self.filename}")
return self.file
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 FileNotFoundError:
print(f"文件不存在: {exc_val}")
return True # 抑制异常
return False # 其他异常继续抛出
# 使用示例
with FileManager('test.txt', 'w') as f:
f.write('Hello, Context Manager!')
with FileManager('nonexistent.txt') as f:
content = f.read() # 文件不存在的异常被抑制
数据库连接管理
import sqlite3
from contextlib import contextmanager
class DatabaseConnection:
def __init__(self, db_path):
self.db_path = db_path
self.connection = None
def __enter__(self):
self.connection = sqlite3.connect(self.db_path)
print(f"连接到数据库: {self.db_path}")
return self.connection
def __exit__(self, exc_type, exc_val, exc_tb):
if exc_type is None:
# 没有异常,提交事务
self.connection.commit()
else:
# 有异常,回滚事务
self.connection.rollback()
print(f"事务回滚: {exc_val}")
self.connection.close()
print("关闭数据库连接")
# 如果不想让事务异常影响程序,返回True
return False
# 使用示例
with DatabaseConnection('test.db') as conn:
cursor = conn.cursor()
cursor.execute('CREATE TABLE IF NOT EXISTS users (id INT, name TEXT)')
cursor.execute('INSERT INTO users VALUES (1, "Alice")')
# 如果这里抛出异常,会自动回滚
使用 @contextmanager 装饰器
计时器案例
from contextlib import contextmanager
import time
@contextmanager
def timer(message="执行时间"):
"""计时器上下文管理器"""
start = time.time()
try:
yield
finally:
elapsed = time.time() - start
print(f"{message}: {elapsed:.4f}秒")
# 使用示例
with timer("数据计算"):
# 模拟耗时操作
sum(i for i in range(10000000))
# 嵌套使用
with timer("总执行时间"):
with timer("子任务1"):
time.sleep(1)
with timer("子任务2"):
time.sleep(0.5)
资源锁定
import threading
from contextlib import contextmanager
@contextmanager
def lock(lock_obj, timeout=5):
"""带超时的锁管理"""
acquired = lock_obj.acquire(timeout=timeout)
try:
if not acquired:
raise TimeoutError("获取锁超时")
yield
finally:
if acquired:
lock_obj.release()
# 使用示例
shared_data = {"counter": 0}
data_lock = threading.Lock()
def increment():
with lock(data_lock, timeout=2):
shared_data["counter"] += 1
print(f"计数器: {shared_data['counter']}")
实际应用场景
日志记录上下文
from contextlib import contextmanager
import logging
@contextmanager
def log_level(level=logging.DEBUG):
"""临时改变日志级别"""
logger = logging.getLogger()
old_level = logger.level
try:
logger.setLevel(level)
yield
finally:
logger.setLevel(old_level)
# 使用示例
logging.basicConfig(level=logging.INFO)
print("正常日志级别: INFO")
with log_level(logging.DEBUG):
logging.debug("这条DEBUG日志会显示")
logging.info("这条INFO日志也会显示")
print("恢复日志级别后")
logging.debug("这条DEBUG日志不会显示")
临时改变工作目录
import os
from contextlib import contextmanager
@contextmanager
def working_directory(new_dir):
"""临时切换工作目录"""
old_dir = os.getcwd()
try:
os.chdir(new_dir)
print(f"切换到目录: {os.getcwd()}")
yield
finally:
os.chdir(old_dir)
print(f"恢复目录: {os.getcwd()}")
# 使用示例
with working_directory('/tmp'):
# 在/tmp目录下操作
print(f"当前目录: {os.getcwd()}")
with open('temp_file.txt', 'w') as f:
f.write('在临时目录创建的文件')
print(f"回到了原目录: {os.getcwd()}")
异常抑制与清理
from contextlib import contextmanager
import traceback
@contextmanager
def safe_execution(raise_exception=False):
"""安全执行,可以决定是否抛出异常"""
try:
yield
except Exception as e:
print(f"捕获异常: {type(e).__name__}: {e}")
traceback.print_exc()
if raise_exception:
raise
else:
print("异常已被抑制,继续执行")
# 使用示例
with safe_execution(raise_exception=False):
result = 1 / 0 # 被抑制
print("程序继续运行")
with safe_execution(raise_exception=True):
result = 1 / 0 # 会抛出异常
print("这行不会执行")
高级应用:自己实现 with 语句
class MyContextManager:
def __init__(self, data):
self.data = data
def __enter__(self):
print(f"进入上下文管理器,数据: {self.data}")
self.data = self.data.upper()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
print(f"退出上下文管理器")
if exc_type is not None:
print(f"异常类型: {exc_type.__name__}")
print(f"异常信息: {exc_val}")
# 你可以在这里进行自定义的异常处理
return False
def __str__(self):
return f"MyContextManager(data={self.data})"
# 使用示例
with MyContextManager("hello world") as cm:
print(f"上下文中的数据: {cm}")
# 如果这里抛出异常
# raise ValueError("测试异常")
异步上下文管理器(Python 3.7+)
import asyncio
from contextlib import asynccontextmanager
@asynccontextmanager
async def async_resource_manager():
"""异步资源管理器"""
print("异步获取资源")
resource = {"data": "async_data"}
try:
yield resource
finally:
print("异步释放资源")
resource.clear()
# 使用示例
async def main():
async with async_resource_manager() as res:
print(f"使用资源: {res}")
await asyncio.sleep(1)
asyncio.run(main())
实用技巧
# 在一个with语句中管理多个资源
with open('file1.txt', 'w') as f1, open('file2.txt', 'w') as f2:
f1.write('File 1 content')
f2.write('File 2 content')
# 使用contextlib中的其他工具
from contextlib import closing, redirect_stdout, redirect_stderr
# 自动调用close方法
class CustomResource:
def close(self):
print("资源已清理")
with closing(CustomResource()):
pass # 自动调用close()
# 重定向标准输出
with open('output.log', 'w') as f:
with redirect_stdout(f):
print("这条信息会写入文件,而不是控制台")
这些案例涵盖了上下文管理器的主要使用场景:资源管理、异常处理、状态保存恢复等,在实际开发中,上下文管理器特别适合处理需要成对操作的场景(open/close、lock/unlock、begin/rollback-commit等)。
标签: 上下文管理器