Python 上下文管理器

# 传统方式（容易忘记关闭）
f = open('test.txt', 'w')
f.write('hello')
f.close()

# 使用 with（自动关闭）
with open('test.txt', 'w') as f:
    f.write('hello')
# 文件自动关闭，即使发生异常

# 读取一个文件，写入另一个文件
with open('input.txt', 'r') as fin, open('output.txt', 'w') as fout:
    content = fin.read()
    fout.write(content.upper())

class FileManager:
    """文件管理上下文管理器"""
    
    def __init__(self, filename, mode):
        self.filename = filename
        self.mode = mode
        self.file = None
    
    def __enter__(self):
        """进入上下文时调用"""
        print(f"打开文件: {self.filename}")
        self.file = open(self.filename, self.mode)
        return self.file
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """退出上下文时调用"""
        print(f"关闭文件: {self.filename}")
        if self.file:
            self.file.close()
        
        # 返回 True 会吞掉异常，返回 False 会抛出异常
        return False

# 使用
with FileManager('test.txt', 'w') as f:
    f.write('测试内容')

class ExceptionHandler:
    """演示异常处理"""
    
    def __enter__(self):
        print("进入上下文")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        exc_type: 异常类型
        exc_val: 异常值
        exc_tb: 异常追踪信息
        """
        if exc_type:
            print(f"捕获到异常: {exc_type.__name__}: {exc_val}")
            # return True  # 吞掉异常
            return False   # 继续抛出异常
        print("正常退出")
        return True

# 测试正常情况
print("=== 正常情况 ===")
with ExceptionHandler():
    print("执行业务逻辑")

# 测试异常情况
print("\n=== 异常情况 ===")
try:
    with ExceptionHandler():
        print("即将抛出异常")
        raise ValueError("出错了!")
except ValueError:
    print("外部捕获到异常")

from contextlib import contextmanager

@contextmanager
def managed_file(filename, mode):
    """使用装饰器创建上下文管理器"""
    print(f"打开: {filename}")
    file = open(filename, mode)
    try:
        yield file  # yield 之前的代码是 __enter__
    finally:
        print(f"关闭: {filename}")
        file.close()
        # yield 之后的代码是 __exit__

# 使用
with managed_file('test.txt', 'w') as f:
    f.write('使用 contextmanager')

from contextlib import contextmanager

@contextmanager
def timer():
    """计时上下文管理器"""
    import time
    start = time.time()
    yield
    elapsed = time.time() - start
    print(f"耗时: {elapsed:.4f} 秒")

# 使用
with timer():
    total = sum(range(1000000))
    print(f"计算结果: {total}")

from contextlib import contextmanager

class MockDB:
    """模拟数据库"""
    def __init__(self):
        self.data = {}
        self.in_transaction = False
    
    def begin(self):
        self.in_transaction = True
        print("开始事务")
    
    def commit(self):
        self.in_transaction = False
        print("提交事务")
    
    def rollback(self):
        self.in_transaction = False
        print("回滚事务")
    
    def insert(self, key, value):
        if not self.in_transaction:
            raise RuntimeError("不在事务中")
        self.data[key] = value
        print(f"插入: {key} = {value}")

db = MockDB()

@contextmanager
def transaction(database):
    """事务上下文管理器"""
    database.begin()
    try:
        yield database
        database.commit()
    except Exception as e:
        database.rollback()
        raise e

# 使用事务
with transaction(db) as conn:
    conn.insert("user1", "Alice")
    conn.insert("user2", "Bob")

from contextlib import suppress

# 传统方式
try:
    os.remove('不存在的文件.txt')
except FileNotFoundError:
    pass

# 使用 suppress
with suppress(FileNotFoundError):
    import os
    os.remove('不存在的文件.txt')

from contextlib import redirect_stdout
import io

# 捕获 print 输出
output = io.StringIO()
with redirect_stdout(output):
    print("这条信息被捕获了")
    print("这也是")

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

from contextlib import ExitStack

# 需要管理多个不确定数量的资源
files = ['file1.txt', 'file2.txt', 'file3.txt']

with ExitStack() as stack:
    opened_files = [
        stack.enter_context(open(f, 'w'))
        for f in files
    ]
    
    for i, f in enumerate(opened_files):
        f.write(f"内容 {i + 1}")
    
    # 所有文件会自动关闭

import os
from contextlib import contextmanager

@contextmanager
def temporary_env(**kwargs):
    """临时修改环境变量"""
    old_values = {}
    for key, value in kwargs.items():
        old_values[key] = os.environ.get(key)
        os.environ[key] = value
    
    try:
        yield
    finally:
        for key, old_value in old_values.items():
            if old_value is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = old_value

# 使用
print(f"当前 PATH: {os.environ.get('MY_VAR', '未设置')}")

with temporary_env(MY_VAR="临时值"):
    print(f"临时 PATH: {os.environ.get('MY_VAR')}")

print(f"恢复后 PATH: {os.environ.get('MY_VAR', '未设置')}")

import os
from contextlib import contextmanager

@contextmanager
def change_directory(path):
    """临时切换工作目录"""
    original = os.getcwd()
    os.chdir(path)
    try:
        yield original
    finally:
        os.chdir(original)

# 使用
print(f"当前目录: {os.getcwd()}")
# with change_directory('/tmp'):
#     print(f"临时目录: {os.getcwd()}")
# print(f"恢复目录: {os.getcwd()}")

from threading import Lock
from contextlib import contextmanager

lock = Lock()

@contextmanager
def acquire_lock(lock_obj):
    """自动获取和释放锁"""
    lock_obj.acquire()
    try:
        yield
    finally:
        lock_obj.release()

# 使用
with acquire_lock(lock):
    print("持有锁，执行临界区代码")

from contextlib import asynccontextmanager

@asynccontextmanager
async def async_resource():
    """异步上下文管理器"""
    print("获取异步资源")
    await asyncio.sleep(0.1)
    try:
        yield "资源"
    finally:
        print("释放异步资源")

# 使用
import asyncio

async def main():
    async with async_resource() as res:
        print(f"使用 {res}")

# asyncio.run(main())