How to create a context manager which is NOT a decorator?

177 views Asked by At

I have a function which looks something like this:

import contextlib

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

It is appropriate for this to be used as a context manager, like this:

with special_context(...):
    ...

... but it is not appropriate for it to be used as a decorator:

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

I understand that Python 3.2 added decorator support to contextlib.contextmanager, but in my API it indicates a mistake which causes bugs. I like the ergonomics of contextlib.contextmanager, but I would like to prevent the API from being misused.

Is there a similar construct available (ideally in the standard libs) which would make special_context a context manager, but not a decorator?

Specifically, I want something like this:

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

Please help me to find or define contextmanager_without_decorator.

3

There are 3 answers

3
dROOOze On

contextlib.ContextDecorator is the part of the return value of @contextlib.contextmanager which grants the ability to redecorate (in your case, @special_context(...)). The simplest way to prevent this from happening at runtime is to disable contextlib.ContextDecorator.__call__. The following is one possible implementation of @contextmanager_without_decorator, reproduced by tracing the runtime implementation of @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
    ...
1
meshy On

Based on the code in contextlib, I've come up with this solution:

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

The version in contextlib has more errorhandling around next, which perhaps this could benefit from, but this is a pleasing foundation.

Decoration now fails with an error to indicate the user's mistake:

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

0
sytech On

Probably to not use contextlib at all and just write your context manager as a class with __enter__ and __exit__ methods, rather than a function.

For example, your existing function context manager

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

Can be rewritten with this usage like:

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

You could also modify the __init__ to fail early if it indicates usage as a function decorator... assuming in practice that your context manager doesn't accept a function as a single positional parameter.

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
    # ...

Then usage as a decorator without arguments fails at definition time, rather than when the function is called:

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

You can also consider adding a __call__ method that provides a more understandable error message like:

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

To get the usage you desire, in a generalized form, you can write a small helper function, as you described in your own answer.