Python中@contextmanager上下文管理器

Flask 中的 @contextmanager 主要用于资源管理，例如数据库连接、文件操作、事务处理等。它并不是 Flask 内置的装饰器，而是 Python contextlib 模块提供的一个上下文管理器工具。

---

1. @contextmanager 基础

在 Python 中，contextlib.contextmanager 允许你使用 with 语句管理资源，而无需创建完整的类。例如：

from contextlib import contextmanager

@contextmanager
def my_resource():
    print("资源初始化")
    yield "资源对象"
    print("资源清理")

# 使用 `with` 语句
with my_resource() as res:
    print(f"正在使用: {res}")

输出：

资源初始化
正在使用: 资源对象
资源清理

其中：

• yield 之前的代码在进入上下文时执行（资源初始化）。
• yield 之后的代码在退出上下文时执行（资源清理）。
• yield 语句的返回值赋给 with 语句中的 as 变量。

---

2. Flask 中的 @contextmanager 使用场景

Flask 主要在数据库连接、会话管理、请求生命周期等场景中使用 @contextmanager，以下是几个常见示例：

2.1 管理数据库连接

假设 Flask 使用 SQLAlchemy 进行数据库操作，每次请求都需要开启和关闭数据库连接：

from flask import Flask
from contextlib import contextmanager
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

app = Flask(__name__)

# 配置数据库
DATABASE_URL = "sqlite:///example.db"
engine = create_engine(DATABASE_URL)
Session = sessionmaker(bind=engine)

@contextmanager
def get_db_session():
    """数据库会话管理器"""
    session = Session()
    try:
        yield session  # 这里返回 session 对象
        session.commit()  # 事务提交
    except Exception:
        session.rollback()  # 发生异常时回滚
        raise
    finally:
        session.close()  # 关闭数据库连接

@app.route("/")
def index():
    with get_db_session() as session:
        result = session.execute("SELECT 'Hello Flask!'").fetchone()
        return result[0]

if __name__ == "__main__":
    app.run(debug=True)

执行流程：

1. get_db_session() 创建数据库会话 session 并返回。

2. 进入 with 代码块时，可以使用 session 进行查询。

3. 代码块执行完毕后，

   yield
   

   后的代码执行：

   • 提交事务（commit()）。
   • 关闭会话（close()）。
   • 若发生异常，则回滚事务（rollback()）。

---

2.2 Flask app.app_context() 手动管理应用上下文

在 Flask 运行时，某些操作需要访问 app 的应用上下文（Application Context），例如 current_app 或 g 对象。但是在某些情况下（如脚本运行、后台任务），可能需要手动管理上下文：

@contextmanager
def app_context():
    """手动管理 Flask 上下文"""
    ctx = app.app_context()
    ctx.push()  # 进入应用上下文
    try:
        yield
    finally:
        ctx.pop()  # 退出应用上下文

with app_context():
    print("在 Flask 应用上下文中执行代码")

这样可以确保代码在 Flask 上下文环境下执行，否则会报 RuntimeError: Working outside of application context.

---

2.3 Flask request_context() 手动管理请求上下文

类似地，Flask 也允许手动管理请求上下文，用于在后台任务或测试中模拟 HTTP 请求环境：

@contextmanager
def request_context(path="/"):
    """手动创建请求上下文"""
    ctx = app.test_request_context(path)
    ctx.push()  # 进入请求上下文
    try:
        yield
    finally:
        ctx.pop()  # 退出请求上下文

with request_context("/hello"):
    from flask import request
    print(f"当前请求路径: {request.path}")  # /hello

---

2.4 处理临时文件

使用 @contextmanager 方便管理临时文件，确保使用后自动删除：

import os
import tempfile
from contextlib import contextmanager

@contextmanager
def temporary_file(suffix=".txt"):
    """创建临时文件并自动删除"""
    temp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
    try:
        yield temp.name  # 传出文件路径
    finally:
        os.remove(temp.name)  # 退出时删除文件

@app.route('/generate')
def generate_file():
    with temporary_file() as file_path:
        with open(file_path, "w") as f:
            f.write("Hello Flask!")
        return f"临时文件路径: {file_path}"

---

3. 为什么使用 @contextmanager

相较于手动管理资源，@contextmanager 让代码更简洁、优雅，避免资源泄漏：

不使用 @contextmanager

def db_operation():
    session = Session()
    try:
        result = session.query(User).all()
        session.commit()
        return result
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()

使用 @contextmanager

with get_db_session() as session:
    result = session.query(User).all()

✅ 代码简洁
✅ 资源自动管理
✅ 避免手动 try-finally

---

总结

1. @contextmanager 是 Python 提供的上下文管理工具，用于管理资源的创建与释放。
2. 在 Flask 中常用于：
   • 数据库会话管理
   • 手动管理 Flask 上下文
   • 管理请求上下文
   • 文件操作
3. 使用 @contextmanager 简化代码，提高可读性，确保资源自动释放，防止泄漏。

你可以试试看如何在你的 Flask 项目中优化资源管理 😉