如何编写一致的有状态上下文管理器？

Question

编辑：正如 Thierry Lathuille, PEP567, where ContextVar was introduced, was not designed to address generators (unlike the withdrawn PEP550 指出的那样）。尽管如此，主要问题仍然存在。如何编写在多线程、生成器和 asyncio 任务中正确运行的有状态上下文管理器？

我有一个库，其中包含一些可以在不同 "modes" 中工作的函数，因此它们的行为可以通过本地上下文进行更改。我正在查看 contextvars 模块来可靠地实现它，因此我可以从不同的线程、异步上下文等使用它。但是，我无法让一个简单的示例正常工作。考虑这个最小的设置：

from contextlib import contextmanager
from contextvars import ContextVar

MODE = ContextVar('mode', default=0)

@contextmanager
def use_mode(mode):
    t = MODE.set(mode)
    try:
        yield
    finally:
        MODE.reset(t)

def print_mode():
   print(f'Mode {MODE.get()}')

这是一个带有生成器函数的小测试：

def first():
    print('Start first')
    print_mode()
    with use_mode(1):
        print('In first: with use_mode(1)')
        print('In first: start second')
        it = second()
        next(it)
        print('In first: back from second')
        print_mode()
        print('In first: continue second')
        next(it, None)
        print('In first: finish')

def second():
    print('Start second')
    print_mode()
    with use_mode(2):
        print('In second: with use_mode(2)')
        print('In second: yield')
        yield
        print('In second: continue')
        print_mode()
        print('In second: finish')

first()

我得到以下输出：

Start first
Mode 0
In first: with use_mode(1)
In first: start second
Start second
Mode 1
In second: with use_mode(2)
In second: yield
In first: back from second
Mode 2
In first: continue second
In second: continue
Mode 2
In second: finish
In first: finish

在部分中：

In first: back from second
Mode 2
In first: continue second

它应该是 Mode 1 而不是 Mode 2，因为这是从 first 打印的，据我所知，应用上下文应该是 use_mode(1)。但是，似乎 second 的 use_mode(2) 堆叠在它上面，直到生成器完成。 contextvars 不支持生成器吗？如果是这样，是否有任何方法可以可靠地支持有状态的上下文管理器？可靠地，我的意思是无论我是否使用它都应该表现一致：

多线程。
发电机。
asyncio

Answer 1

你实际上在那里有一个“互锁上下文”——如果不返回 second 函数的 __exit__ 部分，它不会恢复上下文用ContextVars，无论如何。

所以，我想出了一些办法——而且是我能想到的最好的办法是一个装饰器，用于显式声明哪个可调用对象将拥有自己的上下文 - 我创建了一个 ContextLocal class 用作命名空间，就像 thread.local - 并且该命名空间中的属性应该按照您的预期正常运行。

我现在正在完成代码 - 所以我还没有针对 async 或多线程测试它，但它应该可以工作。如果你能帮我写一个合适的测试，下面的解决方案本身就可以成为一个 Python 包。

（我不得不求助于在生成器和协同例程框架局部字典中注入一个对象，以便在生成器或协同例程结束后清理上下文注册表 - PEP 558 将locals() 对于 Python 3.8+ 的行为，我现在不记得这种注入是否被允许了——不过它可以达到 3.8 beta 3，所以我认为这种用法是有效的）。

无论如何，这是代码（命名为 context_wrapper.py）：

"""
Super context wrapper -

meant to be simpler to use and work in more scenarios than
Python's contextvars.

Usage:
Create one or more project-wide instances of "ContextLocal"
Decorate your functions, co-routines, worker-methods and generators
that should hold their own states with that instance's `context` method -

and use the instance as namespace for private variables that will be local
and non-local until entering another callable decorated
with `intance.context` - that will create a new, separated scope
visible inside  the decorated callable.


"""

import sys
from functools import wraps

__author__ = "João S. O. Bueno"
__license__ = "LGPL v. 3.0+"

class ContextError(AttributeError):
    pass


class ContextSentinel:
    def __init__(self, registry, key):
        self.registry = registry
        self.key = key

    def __del__(self):
        del self.registry[self.key]


_sentinel = object()


class ContextLocal:

    def __init__(self):
        super().__setattr__("_registry", {})

    def _introspect_registry(self, name=None):

        f = sys._getframe(2)
        while f:
            h = hash(f)
            if h in self._registry and (name is None or name in self._registry[h]):
                return self._registry[h]
            f = f.f_back
        if name:
            raise ContextError(f"{name !r} not defined in any previous context")
        raise ContextError("No previous context set")


    def __getattr__(self, name):
        namespace = self._introspect_registry(name)
        return namespace[name]


    def __setattr__(self, name, value):
        namespace = self._introspect_registry()
        namespace[name] = value


    def __delattr__(self, name):
        namespace = self._introspect_registry(name)
        del namespace[name]

    def context(self, callable_):
        @wraps(callable_)
        def wrapper(*args, **kw):
            f = sys._getframe()
            self._registry[hash(f)] = {}
            result = _sentinel
            try:
                result = callable_(*args, **kw)
            finally:
                del self._registry[hash(f)]
                # Setup context for generator or coroutine if one was returned:
                if result is not _sentinel:
                    frame = getattr(result, "gi_frame", getattr(result, "cr_frame", None))
                    if frame:
                        self._registry[hash(frame)] = {}
                        frame.f_locals["$context_sentinel"] = ContextSentinel(self._registry, hash(frame))

            return result
        return wrapper

这里是你的例子的修改版本：

from contextlib import contextmanager

from context_wrapper import ContextLocal

ctx = ContextLocal()


@contextmanager
def use_mode(mode):
    ctx.MODE = mode
    print("entering use_mode")
    print_mode()
    try:
        yield
    finally:

        pass

def print_mode():
   print(f'Mode {ctx.MODE}')


@ctx.context
def first():
    ctx.MODE = 0
    print('Start first')
    print_mode()
    with use_mode(1):
        print('In first: with use_mode(1)')
        print('In first: start second')
        it = second()
        next(it)
        print('In first: back from second')
        print_mode()
        print('In first: continue second')
        next(it, None)
        print('In first: finish')
        print_mode()
    print("at end")
    print_mode()

@ctx.context
def second():
    print('Start second')
    print_mode()
    with use_mode(2):
        print('In second: with use_mode(2)')
        print('In second: yield')
        yield
        print('In second: continue')
        print_mode()
        print('In second: finish')

first()

这是运行的输出：

Start first
Mode 0
entering use_mode
Mode 1
In first: with use_mode(1)
In first: start second
Start second
Mode 1
entering use_mode
Mode 2
In second: with use_mode(2)
In second: yield
In first: back from second
Mode 1
In first: continue second
In second: continue
Mode 2
In second: finish
In first: finish
Mode 1
at end
Mode 1

（它将比原生上下文变量慢几个数量级是内置的 Python 运行时本机代码 - 但看起来等量使用更容易全神贯注）

如何编写一致的有状态上下文管理器？

How do I write consistent stateful context managers?

python

python-asyncio

python-contextvars