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

Question

编辑:正如 Thierry Lathuille 所指出的那样, PEP567 ，在引入 ContextVar 的地方，并不是为了解决生成器问题而设计的(与撤回的 PEP550 不同)。尽管如此，主要问题仍然存在。如何编写在多线程、生成器和 asyncio 任务中正确运行的有状态上下文管理器？

---

我有一个库，其中包含一些可以在不同“模式”下工作的函数，因此它们的行为可以通过本地上下文进行更改。我在看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 不支持生成器吗？如果是这样，是否有任何方法可以可靠地支持有状态的上下文管理器？我所说的可靠是指无论我是否使用它都应该表现一致:

• 多线程。
• 生成器。
• 异步

Answer 1

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

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

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

(我不得不求助于在生成器和协同例程框架局部字典中注入(inject)一个对象，以便在生成器或协同例程结束后清理上下文注册表 - 有 PEP 558 形式化 locals() 用于 Python 3.8+，我现在不记得这种注入(inject)是否被允许了——尽管它可以工作到 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 运行时 native 代码 - 但看起来 更容易全神贯注地使用相同的数量)