如何创建一个不是装饰器的上下文管理器？

Question

我有一个函数，看起来像这样：

import contextlib

@contextlib.contextmanager
def special_context(...):
    ...
    yield
    ...

这个函数适合用作上下文管理器，像这样使用：

with special_context(...):
    ...

但是，它不适合用作装饰器：

# Not OK:
@special_context(...)
def foo():
    ...

我知道Python 3.2增加了对contextlib.contextmanager的装饰器支持，但在我的API中，这会导致错误，进而引发bug。我喜欢contextlib.contextmanager的使用体验，但我想防止API被错误使用。

有没有类似的结构（最好是在标准库中），可以让special_context成为上下文管理器，但不作为装饰器使用？

具体来说，我想要这样的东西：

@contextmanager_without_decorator
def special_context(...):
    ...
    yield
    ...

请帮我找到或定义contextmanager_without_decorator。

Answer 1

根据这个在contextlib里的代码，我想出了这个解决方案：

def contextmanager_without_decorator(wrapped):
    class _ContextManager:

        def __init__(self, *args, **kwargs):
            self.args = args
            self.kwargs = kwargs

        def __enter__(self):
            self.gen = wrapped(*self.args, **self.kwargs)
            next(self.gen)

        def __exit__(self, *args):
            next(self.gen, None)

    return _ContextManager

在contextlib里的版本对next有更多的错误处理，或许这个解决方案也可以借鉴，但这已经是一个不错的基础了。

现在，如果装饰器出错，会给出一个错误提示，告诉用户哪里出错了：

@special_context(...)
def foo():  # TypeError: '_ContextManager' object is not callable
    ...

Answer 2

contextlib.ContextDecorator 是 @contextlib.contextmanager 返回值的一部分，它让你可以重新装饰（在你的例子中，就是 @special_context(...)）。为了防止在运行时发生这种情况，最简单的方法就是禁用 contextlib.ContextDecorator.__call__。下面是 @contextmanager_without_decorator 的一种可能实现方式，这个实现是通过跟踪 @contextlib.contextmanager 的运行时实现得来的：

from __future__ import annotations

import typing as t

import contextlib
import functools

if t.TYPE_CHECKING:
    import collections.abc as cx
    import typing_extensions as tx

    _P = tx.ParamSpec("_P")
    _ContextManagerDecoratee: tx.TypeAlias = cx.Callable[_P, cx.Generator["_T_co", None, None]]

_T_co = t.TypeVar("_T_co", covariant=True)

class _GeneratorContextManager(contextlib._GeneratorContextManager[_T_co]):
    __call__: t.ClassVar[None] = None  # type: ignore[assignment]

def contextmanager_without_decorator(func: _ContextManagerDecoratee[_P, _T_co], /) -> cx.Callable[_P, _GeneratorContextManager[_T_co]]:
    @functools.wraps(func)
    def helper(*args: _P.args, **kwds: _P.kwargs) -> _GeneratorContextManager[_T_co]:
        return _GeneratorContextManager(func, args, kwds)

    return helper

@contextmanager_without_decorator
def special_context() -> cx.Generator[int, None, None]:
    yield 1

with special_context() as num:
    assert num == 1  # OK

# Re-decoration disabled
@special_context()
def foo():  # TypeError: 'NoneType' object is not callable
    ...

Answer 3

可能不需要使用 contextlib，而是直接把你的上下文管理器写成一个类，里面包含 __enter__ 和 __exit__ 这两个方法，而不是用一个函数。

比如，你现在的函数上下文管理器

@contextmanager
def special_context(*args, **kwargs):
    ...
    yield 1
    ...

可以用这种方式重写：

class special_context:
    def __init__(self, *args, **kwargs):
        self.args = args
        self.kwargs = kwargs

    def __enter__(self):
        ...
        return 1

    def __exit__(self, exc_type, exc_val, exc_tb):
        ...

with special_context(...) as ctx:
    ... # works

@special_context(...)
def foo():
   ... # fails

@special_context
def foo():
    ...

foo() # fails

你还可以修改 __init__ 方法，让它在使用时如果被当作函数装饰器来用就提前报错……假设实际上你的上下文管理器不接受一个函数作为位置参数。

class special_context:
    def __init__(self, *args, **kwargs):
        if len(args) == 1 and not kwargs:
            if inspect.isfunction(args[0]):
                raise ValueError(f"You can't use {self.__class__.__name__} as a decorator")
        # ... rest as before
    # ...

这样在定义的时候，如果当作装饰器使用但没有参数，就会报错，而不是在函数被调用时才报错：

# fails at definition time with:
# ValueError: You can't use special_context as a decorator
@special_context
def foo(): 
    ...

你还可以考虑添加一个 __call__ 方法，这样可以提供一个更易懂的错误信息，比如：

def __call__(self, *args, **kwargs):
    raise ValueError(f"{self.__class__.__name__} cannot be used as a decorator")

为了达到你想要的用法，可以写一个小的辅助函数，就像你在你的 回答 中描述的那样。