Python上下文管理-contextvars模块详解

我爱让机器学习

已于 2024-09-24 22:31:50 修改

阅读量940

点赞数 30

文章标签： python contextvars 异步 asyncio

于 2024-09-24 18:21:34 首次发布

本文链接：https://blog.csdn.net/youngwyj/article/details/142494422

版权

什么是 `contextvars`？

contextvars 是 Python 中的一个模块，用于处理上下文变量，这些变量支持对异步任务（如协程）中的数据进行隔离和局部存储。上下文变量让我们能够在异步任务或线程中保持变量的独立性，从而避免数据泄露或干扰。

想象一下，你在一家饭店工作，每位客人都有自己的一张账单。在传统的编程中，所有客人共享同一个账单（即全局变量），这很容易出错。而 contextvars 就像是给每位客人发了一张个人账单，确保他们点的菜对不上其他人的账单。

基本概念

上下文变量：
- 上下文变量是 ContextVar 类的实例。它们用于保存和管理与特定执行上下文关联的数据。
- 这些数据的生命周期与其所在的上下文一致，一旦上下文结束，数据也就不复存在。
上下文：
- Context 是一组上下文变量及其值的集合。你可以将上下文理解为一个特殊的“数据存储区”，在特定的异步执行流中传递和存储数据。

主要组件

ContextVar：
- 创建上下文变量。例如，var = ContextVar('var_name')。
- 提供了 get() 和 set() 方法用于访问和修改变量的值。
Token：
- ContextVar.set() 方法返回一个 Token，表示更改前的变量状态。你可以通过 var.reset(token) 方法恢复变量的旧状态。
Context：
- 表示一组上下文变量的完整状态，可以手动管理上下文数据。
- 通过 copy_context() 函数可以复制当前上下文，用于在新的执行流中共享或继续使用当前上下文的状态。

使用场景

contextvars 模块非常有用的场景包括：

异步编程：在异步任务中使用上下文变量可以确保在一个任务中设置的变量值不会影响其他任务。
处理请求：在请求生命周期中存储和访问相关数据，而无需显式传递参数。
全局状态管理：替代全局变量以避免数据在不同执行流间的意外共享。

创建和使用 ContextVar

首先，我们来看如何创建和使用一个基本的上下文变量：

import contextvars

# 创建一个上下文变量，相当于给每位客人准备一个独立的账单
my_var = contextvars.ContextVar('my_var')

# 设置初始值（可以没有）
my_var.set('666')

# 读取值
print('Initial value:', my_var.get())


"""
输出: 
Initial value: 666
"""

异步环境中的使用

使用协程（coroutine）时,contextvars 确保在并发执行时，每个协程都有独立的数据。例如：

import contextvars
import asyncio

# 创建上下文变量
var = contextvars.ContextVar('var', default='no_value')


async def worker(worker_id):
    print(f'Worker {worker_id} 设置前的值: {var.get()}')  # 打印默认值
    # 这里，给上下文变量设置不同的值
    var.set(f'value_for_worker_{worker_id}')
    print(f'Worker {worker_id} 设置后的值: {var.get()}')  # 每个 worker 获取到的值是独立的
    await asyncio.sleep(1)  # 模拟其他操作
    print(f'Worker {worker_id} 结束时的值: {var.get()}')  # 再次打印以确认变量值保持不变


async def main():
    async with asyncio.TaskGroup() as tg:
        tg.create_task(worker(1))
        tg.create_task(worker(2))


asyncio.run(main())


"""
输出：
Worker 1 设置前的值: no_value
Worker 1 设置后的值: value_for_worker_1
Worker 2 设置前的值: no_value
Worker 2 设置后的值: value_for_worker_2
Worker 1 结束时的值: value_for_worker_1
Worker 2 结束时的值: value_for_worker_2
"""

解释：

每个 worker 相当于一个客人，他们同时点菜，并记录在各自的“账单”上（var 的值）。
contextvars 保证即便它们几乎同时执行，var 的值也不会相互干扰，保证他们的独立性。

利用 Token 进行状态恢复

在饭店，客人可能会因为预算恢复上一个选项。contextvars 提供了这种回退功能：

import contextvars
import asyncio

# 创建一个上下文变量
var = contextvars.ContextVar('var', default='default_value')


async def worker(name):
    token = var.set(f'{name}_new')
    print(f'{name} set: {var.get()}')  # 设定并打印新值

    # 进行某种操作，随后恢复原状态
    await asyncio.sleep(1)
    var.reset(token)
    print(f'{name} reset: {var.get()}')  # 恢复之前状态并打印


async def main():
    async with asyncio.TaskGroup() as tg:
        tg.create_task(worker(1))
        tg.create_task(worker(2))


asyncio.run(main())


"""
输出：
1 set: 1_new
2 set: 2_new
1 reset: default_value
2 reset: default_value
"""

解释：

token 是一种“撤销”工具，允许我们回到某个之前的状态。
每次 set 返回的 token 代表该操作之前的状态。
在需要恢复到某个状态时，调用 var.reset(token)。

在多上下文情况下使用

有时您可能会手动管理和复制上下文，以便它在多个执行流中共享。此时可以使用 Context 对象：

import contextvars

# 创建上下文变量
var = contextvars.ContextVar('var', default='default_value')

# 在默认上下文中设置一个新值
var.set('new_value')

# 使用 copy_context 来复制当前上下文
ctx = contextvars.copy_context()

# 修改原上下文中的变量
var.set('another_value')

# 在复制的上下文中进行操作，验证值不变
print('原上下文中的值:', var.get())  # 打印当前上下文值
print('复制的上下文的值:', ctx[var])  # 打印复制上下文的值（旧值）


"""
输出：
原上下文中的值: another_value
复制的上下文的值: new_value
"""

解释：

copy_context() 方法创建了一份当前上下文的快照。
在原始上下文中修改后，复制的上下文值保持不变，展示了上下文的独立性。

“饭店上下文管理”示例

首先，我们定义一个用于管理“客人账单”的数据结构和相关的上下文管理函数。

from contextvars import ContextVar
from dataclasses import dataclass, field


@dataclass
class GuestBill:
    guest_id: str
    items: dict[str, float] = field(default_factory=dict)  # 记录点的菜品和价格
    table_number: int | None = None

    def add_item(self, item_name: str, price: float) -> None:
        """添加菜品到账单"""
        self.items[item_name] = price

    def remove_item(self, item_name: str) -> None:
        """从账单移除菜品"""
        if item_name in self.items:
            del self.items[item_name]

    def total(self) -> float:
        """计算账单总价"""
        return sum(self.items.values())

    def __repr__(self) -> str:
        return f"GuestBill(guest_id={self.guest_id}, items={self.items}, table={self.table_number})"


_guest_bill_context: ContextVar[GuestBill | None] = ContextVar(
    "GuestBillContext", default=None
)


def current_bill() -> GuestBill | None:
    """获取当前客人的账单"""
    return _guest_bill_context.get()


def ensure_bill() -> GuestBill:
    """确保当前上下文有账单，否则抛出异常"""
    bill = current_bill()
    if bill is None:
        raise RuntimeError("No guest bill context available")
    return bill


def set_bill(bill: GuestBill) -> None:
    """为当前上下文设置账单"""
    _guest_bill_context.set(bill)


def reset_bill() -> None:
    """重置当前上下文的账单"""
    _guest_bill_context.set(None)

接下来，我们模拟一个简单的场景：每位客人可以点菜、撤销菜品，以及获取账单总价。

async def process_guest(guest_id: str, table_number: int):
    # 初始化客人账单上下文
    bill = GuestBill(guest_id=guest_id, table_number=table_number)
    set_bill(bill)

    # 客人开始点菜
    print(f"{guest_id} 初始账单: {current_bill()}")
    current_bill().add_item('Spaghetti', 12.99)
    await asyncio.sleep(1)  # 模拟异步等待

    current_bill().add_item('Salad', 7.99)
    print(f"{guest_id} 加入菜品后账单: {current_bill()}")
    await asyncio.sleep(1)  # 模拟异步等待

    # 撤销一个菜品
    current_bill().remove_item('Salad')
    print(f"{guest_id} 移除一个菜品后账单: {current_bill()}")

    # 获取总价
    total_price = current_bill().total()
    print(f"{guest_id} 账单总价: ${total_price:.2f}")

    # 结账后清除上下文
    reset_bill()

模拟多个客人同时处理

async def main():
    import asyncio

    # 模拟三个客人同时处理
    async with asyncio.TaskGroup() as tg:
        tg.create_task(process_guest('张三', 666))
        tg.create_task(process_guest('李四', 888))
        tg.create_task(process_guest('王麻子', 999))

# 运行异步主函数
asyncio.run(main())


"""
输出：
张三 初始账单: GuestBill(guest_id=张三, items={}, table=666)
李四 初始账单: GuestBill(guest_id=李四, items={}, table=888)
王麻子 初始账单: GuestBill(guest_id=王麻子, items={}, table=999)
张三 加入菜品后账单: GuestBill(guest_id=张三, items={'Spaghetti': 12.99, 'Salad': 7.99}, table=666)
李四 加入菜品后账单: GuestBill(guest_id=李四, items={'Spaghetti': 12.99, 'Salad': 7.99}, table=888)
王麻子 加入菜品后账单: GuestBill(guest_id=王麻子, items={'Spaghetti': 12.99, 'Salad': 7.99}, table=999)
张三 移除一个菜品后账单: GuestBill(guest_id=张三, items={'Spaghetti': 12.99}, table=666)
张三 账单总价: $12.99
李四 移除一个菜品后账单: GuestBill(guest_id=李四, items={'Spaghetti': 12.99}, table=888)
李四 账单总价: $12.99
王麻子 移除一个菜品后账单: GuestBill(guest_id=王麻子, items={'Spaghetti': 12.99}, table=999)
王麻子 账单总价: $12.99
"""

解释:

数据类 GuestBill：用于表示每个客人的账单，包含菜品和价格。
ContextVar 使用：_guest_bill_context 用于存储当前客人的账单上下文。
管理函数：
- current_bill()：获取当前客人账单。
- ensure_bill()：确保账单存在，否则抛出异常。
- set_bill()：为当前客人设置账单。
- reset_bill()：清除当前上下文中的账单。
异步函数 process_guest：
- 使用 await asyncio.sleep() 模拟异步操作之间的等待。
- 每个 process_guest 调用都是一个独立的协程，模拟客人异步点菜的过程。
TaskGroup+create_task：
- 用于并发处理多个客人的账单，确保每个协程任务都有独立的上下文。

通过这种方式，我们确保了即便并发处理多个客人时（并发任务），每个客人（任务）都有独立的账单，极大地减少了数据干扰和可能的错误情况。

注意事项

上下文变量应该在模块顶层创建，而不是在闭包中，以确保它们可以被正确地垃圾回收：
- 当上下文变量在模块顶层创建时，其生命周期与模块绑定，通常会持续到程序结束。这可以确保上下文变量在程序运行期间始终可用，并且避免了意外的垃圾回收问题。
- 在模块顶层创建上下文变量，可以确保它们在整个模块内可见和可访问。这在大型代码库中尤为重要，避免了在局部作用域中重新创建相同名称的上下文变量的问题。
- 如果上下文变量在闭包或局部作用域中创建，那么一旦闭包或作用域结束，该变量可能会被垃圾回收。这会导致丢失上下文变量的状态，从而引发不可预测的行为。
理解 Context 和 ContextVar 的区别：
- Context: 是上下文对象，保存一组特定的上下文变量及其值。
- ContextVar: 是上下文变量对象，用来定义和获取上下文变量的值。
正确使用 ContextVar 的 set 方法:
- set 方法会返回一个 Token，可以用来恢复之前的值，但要小心 Token 的作用域和生命周期，避免滥用导致不可预测的行为。
避免在多线程或异步任务中共享 ContextVar 的值：
- ContextVar 的值是线程和异步任务特定的，如果你在一个异步任务中修改了某个 ContextVar 的值，其他任务不会被影响。
小心上下文变量的嵌套操作：
- 当在嵌套的异步任务中使用 ContextVar 时，需要特别注意变量的赋值和获取，可能会由于不正确的嵌套导致意想不到的值。
理解 copy_context 的作用：
- copy_context 方法用来复制当前的上下文（浅拷贝），这对于将当前上下文传递给新的线程或子任务是非常有用的。