Python 上下文管理器深度实战：从 with 语句到自定义资源管控

1
import sqlite3
2
from contextlib import contextmanager
3

4
@contextmanager
5
def get_db_connection(db_path):
6
    """获取数据库连接，自动提交或回滚事务"""
7
    conn = sqlite3.connect(db_path)
8
    try:
9
        yield conn
10
        conn.commit()  # 正常退出 → 提交
11
    except Exception:
12
        conn.rollback()  # 异常 → 回滚
13
        raise  # 重新抛出，让调用方知道出错了
14
    finally:
15
        conn.close()  # 无论如何都关闭连接
16

17
# 使用
18
with get_db_connection("app.db") as conn:
19
    cursor = conn.cursor()
20
    cursor.execute("CREATE TABLE IF NOT EXISTS users (id INTEGER PRIMARY KEY, name TEXT)")
21
    cursor.execute("INSERT INTO users (name) VALUES (?)", ("Alice",))
22
    # 如果这里抛异常，事务自动回滚；正常结束则自动提交

1
import os
2
from contextlib import contextmanager
3

4
@contextmanager
5
def temp_env(**env_vars):
6
    """临时设置环境变量，退出后自动恢复"""
7
    old_values = {}
8
    for key, value in env_vars.items():
9
        old_values[key] = os.environ.get(key)  # 保存原值（可能是 None）
10
        os.environ[key] = value
11

12
    try:
13
        yield
14
    finally:
15
        for key, old_value in old_values.items():
16
            if old_value is None:
17
                os.environ.pop(key, None)  # 原来没有就删掉
18
            else:
19
                os.environ[key] = old_value  # 恢复原值
20

21
# 使用：测试不同配置
22
with temp_env(DATABASE_URL="sqlite:///test.db", DEBUG="true"):
23
    print(os.environ["DATABASE_URL"])  # sqlite:///test.db
24
    print(os.environ["DEBUG"])         # true
25

26
# 退出后环境变量恢复原状
27
print(os.environ.get("DEBUG"))  # None

1
import fcntl
2
from contextlib import contextmanager
3

4
@contextmanager
5
def file_lock(lock_path):
6
    """基于文件的进程锁"""
7
    lock_file = open(lock_path, 'w')
8
    try:
9
        fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)  # 阻塞式排他锁
10
        print(f"🔐 获取锁: {lock_path}")
11
        yield lock_file
12
    finally:
13
        fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
14
        lock_file.close()
15
        print(f"🔓 释放锁: {lock_path}")
16

17
# 使用
18
with file_lock("/tmp/my_app.lock"):
19
    # 这段代码同一时刻只有一个进程能执行
20
    print("🔧 执行互斥操作...")

1
import os
2
from contextlib import contextmanager
3

4
@contextmanager
5
def working_directory(path):
6
    """临时切换工作目录，退出后恢复"""
7
    original = os.getcwd()
8
    try:
9
        os.chdir(path)
10
        yield path
11
    finally:
12
        os.chdir(original)
13

14
# 使用
15
print(f"当前目录: {os.getcwd()}")
16

17
with working_directory("/tmp"):
18
    print(f"临时目录: {os.getcwd()}")  # /tmp
19
    # 在 /tmp 下执行操作...
20

21
print(f"恢复目录: {os.getcwd()}")  # 回到原目录

1
import io
2
import sys
3
from contextlib import contextmanager
4

5
@contextmanager
6
def capture_output():
7
    """捕获 with 块内的所有 print 输出"""
8
    buffer = io.StringIO()
9
    old_stdout = sys.stdout
10
    sys.stdout = buffer
11
    try:
12
        yield buffer
13
    finally:
14
        sys.stdout = old_stdout
15

16
# 使用
17
with capture_output() as output:
18
    print("这行不会出现在终端")
19
    print("这行也不会")
20

21
captured = output.getvalue()
22
print(f"捕获到 {len(captured)} 个字符")
23
print(f"内容: {captured!r}")

1
from contextlib import ExitStack
2

3
def process_files(file_paths):
4
    """同时打开多个文件，全部处理完再统一关闭"""
5
    with ExitStack() as stack:
6
        files = [
7
            stack.enter_context(open(path, 'r'))
8
            for path in file_paths
9
        ]
10
        # 所有文件都打开了，可以交叉处理
11
        for f in files:
12
            first_line = f.readline().strip()
13
            print(f"{f.name}: {first_line}")
14
    # ExitStack 退出时，所有文件按 LIFO 顺序关闭
15

16
# 使用
17
import tempfile, os
18

19
# 创建测试文件
20
paths = []
21
for i in range(3):
22
    path = os.path.join(tempfile.gettempdir(), f"test_{i}.txt")
23
    with open(path, 'w') as f:
24
        f.write(f"文件 {i} 的内容\n第二行")
25
    paths.append(path)
26

27
process_files(paths)