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

时间:2018-12-04 11:13:24

标签: python

编辑:如Thierry Lathuille所指出,引入ContextVar的{​​{3}}并不是为寻址生成器而设计的(与撤回的PEP567不同)。尽管如此,主要问题仍然存在。如何编写可在多个线程,生成器和asyncio任务下正常运行的有状态上下文管理器?


我有一个库,其中包含一些可以在不同的“模式”下工作的功能,因此可以通过本地上下文来更改其行为。我正在查看PEP550模块,以可靠地实现此功能,因此可以在不同的线程,异步上下文等环境中使用它。但是,我很难获得一个正确的简单示例。考虑以下最小设置:

from contextlib import contextmanager
from contextvars import ContextVar

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

@contextmanager
def use_mode(mode):
    t = MODE.set(mode)
    yield
    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)。但是,似乎use_mode(2)中的second堆叠在其上,直到生成器完成为止。 contextvars不支持生成器吗?如果是这样,有什么方法可以可靠地支持状态上下文管理器?可靠地说,我的意思是无论是否使用它,它的行为都应该一致:

  • 多线程。
  • 发电机。
  • asyncio

1 个答案:

答案 0 :(得分:2)

是的。 棘手的。

您实际上在那里有一个“互锁上下文”-不为__exit__函数返回second部分,它将不会恢复上下文 无论如何,都可以使用ContextVars

所以,我在这里想出了一些东西-我想到的最好的东西 是一个装饰器,用于明确声明哪些可调用对象将具有自己的上下文- 我创建了一个ContextLocal类,该类可用作命名空间,类似于thread.local之类的jsut-并且该命名空间中的属性应按您期望的那样正常运行。

我现在正在整理代码-因此我尚未对其进行异步或多线程测试,但它应该可以工作-如果您可以帮助我编写适当的测试电池,则下面的解决方案本身可能会成为Python软件包。

((我不得不诉诸于在生成器和协同例程框架本地字典中注入一个对象,以便在生成器或协同例程结束后清理上下文注册表-PEP 558形式化了行为适用于Python 3.8及更高版本的locals(),我现在不记得是否允许这种注入-尽管它最多可以使用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

(它将比本地contextvars慢几个数量级 内置在Python运行时本机代码中-但似乎 更容易将它们包装起来以相同数量使用)