Python上下文管理之ContextVar源码解析

测试有道

已于 2023-12-01 12:04:44 修改

阅读量1.1k

点赞数 19

分类专栏： Python 文章标签： python 开发语言

于 2023-11-30 12:16:29 首次发布

本文链接：https://blog.csdn.net/lIujunXHU/article/details/134709067

版权

Python 专栏收录该内容

8 篇文章 0 订阅

订阅专栏

2.1、ContextMeta,ContextVarMeta和TokenMeta三个元类

2.2、contextvar基本的上下文管理方式

2.3、contextvar之Context实现

2.4、Contextvar之Token的实现

2.5、Contextvar之ContextVar的实现

前言

在Python3.7后官方库出现了contextvars模块, 它的主要功能就是可以为多线程以及asyncio生态添加上下文功能,即使程序在多个协程并发运行的情况下,也能调用到程序的上下文变量, 从而使程序逻辑解耦。
在flask框架中，flask1.x版本使用了LocalStack来管理上下文，LocalStack是基于Local即一个字典来实现的上下文管理，在flask2.x版本以后使用了ContextVar管理上下文，那么ContextVar实现上文管理的原理是什么？带个这个问题，分别从ContextVar的使用和源码去看看它是如何完成上下文管理的吧

【注：下面所有Python代码都是基于Python3.9版本的】

【注：如果对异步函数定义和使用不清楚可见】Python异步编程之协程-CSDN博客

1、contextvar的使用

1.1、在普通函数中的使用

1、contextvar的基本使用，包括set()设置值、get()获取值、reset()重置值、获取old_value值、默认值设置

from contextvars import ContextVar

ctx = ContextVar("debug test")
# 设置默认值的方式
# ctx = ContextVar("debug test", default="high")

# 给ctx设置第一个值
token1 = ctx.set("hello")
# 给ctx设置第二个值
token2 = ctx.set("world")

# 获取当前ctx的最新值
print(ctx.get()) # world

# 获取token2对象的旧值
print(token2.old_value) # hello

# 重置token2,此时ctx的当前值变为旧值 hello
ctx.reset(token2)
print(ctx.get()) # hello

# 获取token1的旧值，因为token1没有旧值,所以会返回默认设置对象<Token.MISSING>
print(token1.old_value) # <Token.MISSING>

# 再次把token1的值重置，此时如果ContextVar()初始化没有设置默认值，get()函数返回也没有设置默认值，则会抛出LookupError异常
# 如果存在默认值优先返回get()函数设置的默认值，其次在是ContextVar()初始化设置的默认值
ctx.reset(token1)
print(ctx.get()) # 此处抛出LookupError异常

# get()函数设置默认值方式
# print(ctx.get("tom")) # 此处返回tom

2、run()函数的使用，用于执行一个可调用对象时隔离当前上下文；contextvars.Context 它负责保存 ContextVars 对象和设置的值之间的映射（保存上下文实际是Context完成的），但实际使用中不会直接通过 contextvars.Context 来创建保存上下文，而是通过 contentvars.copy_context 函数来创建。

import contextvars

ctx1 = contextvars.ContextVar("ctx1")
ctx2 = contextvars.ContextVar("ctx2")
ctx1.set("hello")
ctx2.set("world")

# 复制一个当前的上下文对象（注意这是一个带有当前上下文值的新对象）
context = contextvars.copy_context()

# 遍历得到当前上下文所有的ContextVar对象
for ctx, value in context.items():
    print(ctx.get(), ctx.name, value)

    """
    输出结果：
    world ctx2 world
    hello ctx1 hello
    """

def task():
    print(f"befor task: {ctx1.get()} context[ctx1]: {context[ctx1]}")
    ctx1.set("task")
    print(f"after task: {ctx1.get()} context[ctx1]: {context[ctx1]}")

    """
    输出结果：
    befor task: hello context[ctx1]: hello
    after task: task context[ctx1]: task
    """

context.run(task)

print(f"ctx1: {ctx1.get()}  context[ctx1]: {context[ctx1]}") # ctx1: hello  context[ctx1]: task

注意：为什么context[ctx1]在task函数中由hello改变成task，出了task()函数context[ctx1]还是task而不是hello。不应该和ctx1一样保持上下文一致吗？
- ctx1保持上下文一致是因为在task()函数内和在task()函数外使用的上下文不是同一个上下文所以才会保持一致；
- 再看context[ctx1]对象不管是在task()函数内还是在task()函数外context对象自始至终都是同一个，所以context对象中ctx1的值在task()函数内发生了变化，在task()函数外同样也会跟着变化（因为context对象在task()内外指向的都是同一块内存地址）

1.2、在异步函数中的使用

1、异步函数中的使用

在异步函数中是如何实现协程上下文切换呢？我们知道在 asyncio 中，Task 是用于封装协程的对象，它代表一个可以被调度和执行的异步任务；其作用有：
- a.封装协程：Task 对象可以封装一个协程函数或可迭代对象（如生成器），使其变成可调度的异步任务。通过 asyncio.create_task() 或 loop.create_task() 方法创建 Task 对象。
- b.异步任务的调度和执行：Task 对象可以被添加到事件循环中，由事件循环调度和执行。事件循环会在适当的时机切换任务，以实现异步执行。一旦 Task 对象被添加到事件循环中，它就可以与其他任务共享 CPU 时间，实现并发执行。
- c. 取消和异常处理：Task 对象可以被取消，即停止其执行。可以使用 Task.cancel() 方法取消任务，并通过捕获 asyncio.CancelledError 异常来处理取消操作。此外，Task 对象还提供了 add_done_callback() 方法，用于注册回调函数以处理任务执行完毕或抛出异常的情况。
- d. 获取任务状态和结果：通过 Task 对象可以获取任务的当前状态，包括已完成、正在运行、已取消等。使用 Task.done() 方法可以检查任务是否已完成，而 Task.result() 方法可以获取任务的返回值（如果有）或引发的异常（如果有）。

Task初始化源码如下：

class Task(futures._PyFuture):  # Inherit Python Task implementation
                                # from a Python Future implementation.
    _log_destroy_pending = True
    def __init__(self, coro, *, loop=None, name=None):
        super().__init__(loop=loop)
        if self._source_traceback:
            del self._source_traceback[-1]
        if not coroutines.iscoroutine(coro):
            self._log_destroy_pending = False
            raise TypeError(f"a coroutine was expected, got {coro!r}")
        if name is None:
            self._name = f'Task-{_task_name_counter()}'
        else:
            self._name = str(name)

        self._must_cancel = False
        self._fut_waiter = None
        self._coro = coro
        self._context = contextvars.copy_context()
        self._loop.call_soon(self.__step, context=self._context)
        _register_task(self)

从第19行代码可以看到Task中上下文对象是根据contextvars.copy_context()来实现的，这个过程和上面run()函数的使用原理一样，所以我们知道在异步函数中实现上下文管理是通过contextvars.copy_context()来实现的（注：不同Python版本asyncio实现上下文方式可能不一样。contextvars模块是在Python3.7引入的；asyncio 是在 Python 3.4 版本引入的标准库）。
下面是一个示例说明contextvars在异步函数中的使用：

import asyncio
import contextvars

# 创建一个ContextVar
ctx = contextvars.ContextVar("ctx")
ctx.set("hello")

async def task(value, delay):
    print("task: ", ctx.get())
    ctx.set(value)
    await asyncio.sleep(delay)
    print("task: ", ctx.get())

async def main():
    # 异步执行函数
    await asyncio.gather(task('task1', 1), task('task2', 2))
    print("main: ", ctx.get())

asyncio.run(main())

# 执行结果
task:  hello
task:  hello
task:  task1
task:  task2
main:  hello

执行过程：
- 输出结果中的前两个“task: hello”是同时打印出来的，
- 然后先打印出：“task: task1”，在打印出：“task: task2”，
- 最后打印出：“main: hello”。
结论：从上面的执行过程来看异步执行的task('task1', 1),和task('task2', 2)在ctx.set(value)之后内部的上下文是隔离的，执行完成之后又恢复到了之前的上下文。所以在异步函数中使用contextvars模块管理上下文非常方便。

2、contextvar使用的浅拷贝问题

contextvars.copy_context()的作用是拷贝一份当前的上下文对象，注意这里拷贝是浅拷贝，也就是当前上下文的一个引用，只有在发生写数据的时候才会进行深拷贝生成一个新的上下文对象。这里和Linux中的进程fork运用了同一种技术：写时拷贝；这就是为什么我会把上面示例的“在ctx.set(value)之后”标注为红色。下面看一个示例：

import asyncio
import contextvars

ctx1 = contextvars.ContextVar('ctx1')
ctx2 = contextvars.ContextVar("ctx2")
ctx1.set("hello")
ctx2.set({"name": "张三"})
print("befroe ctx1_value_addr: ", id(ctx1.get())) # befroe ctx1_value_addr:  2070012128240
print("befroe ctx2_value_addr: ", id(ctx2.get())) # befroe ctx2_value_addr:  2070012049728

async def main():
    ctx1.set(1)  # 此时会进行写时拷贝
    print("main ctx1_value_addr: ", id(ctx1.get())) # main ctx1_value_addr:  2069998823728
    ctx2_dict = ctx2.get() # 此处只是引用即浅拷贝
    print("main ctx2_value_addr: ", id(ctx2_dict)) # main ctx2_value_addr:  2070012049728
    ctx2_dict["name"] = "李四"
    print(f"main inner ctx1: {ctx1.get()}  ctx2: {ctx2.get()}") # main inner ctx1: 1  ctx2: {'name': '李四'}


print(f"before run ctx1: {ctx1.get()}  ctx2: {ctx2.get()}") # before run ctx1: hello  ctx2: {'name': '张三'}
asyncio.run(main())
print(f"after run  ctx1: {ctx1.get()}  ctx2: {ctx2.get()}") # after run  ctx1: hello  ctx2: {'name': '李四'}

# 执行结果
befroe ctx1_value_addr:  2070012128240
befroe ctx2_value_addr:  2070012049728
before run ctx1: hello  ctx2: {'name': '张三'}
main ctx1_value_addr:  2069998823728
main ctx2_value_addr:  2070012049728
main inner ctx1: 1  ctx2: {'name': '李四'}
after run  ctx1: hello  ctx2: {'name': '李四'}

首先我们分析打印的ctx值，在before run打印的ctx的值和after run的对比发现ctx2的值发生了变化，乍一看显然有点不符合异步执行切换上下文的情况啊，
然后我们再进入到main()的内部看看ctx1和ctx2有什么差别导致出现这样的结果；发现ctx1使用了set()函数也就是重新设置了值，ctx2只使用了get()函数没有重新设置值，结合上面所说的contextvars.copy_context()的作用是拷贝一份当前的上下文对象，只是进行浅拷贝，貌似可以解释这个问题，为了验证实际是不是真是这样，
我们分别在main()外和main()内打印出了ctx1值和ctx2值的地址，发现ctx1值前后的地址发生了改变，ctx2值前后的地址并没有发生变化，正好就对应上面所说的浅拷贝情况。