Improve `contextlib.ContextDecorator` support for generators and async functions · Issue #125862 · python/cpython · GitHub

Suppose we define a context manager, in the usual way:

@contextmanager
def ctx():
    print("enter")
    yield
    print("exit")

def classic():
    with ctx():
        print("body")

classic()  # enter, body, exit

We can also use it as a decorator, thanks to the convenient ContextDecorator class used by @contextmanager:

@ctx()
def fn():
    print("body")

fn()  # enter, body, exit

...but if we naively do the same thing to a generator or an async function, the equivalence breaks down:

@ctx()
def gen():
    print("body")
    yield

for _ in gen(): ...  # enter, exit, body!

@ctx()
async def afn():
    print("body")

await afn()  # enter, exit, body!

This seems pretty obviously undesirable, so I think we'll want to change ContextDecorator.__call__. Possibilities include:

branch on iscoroutinefunction / isgeneratorfunction / isasyncgenfunction, creating alternative inner functions to preserve the invariant that @ctx() is just like writing with ctx(): as the first line of the function body. In a quick survey, this is the expected behavior, but also a change from the current effect of the decorator.
instead of branching, warn (and after a few years raise) when decorating a generator or async function.
- alternative implementation: inspect the return value inside inner, to detect sync wrappers of async functions. I think the increased accuracy is unlikely to be worth the performance cost.
- we could retain the convenience of using a decorator by defining ContextDecorator.gen(), .async_(), and .agen() methods which explicitly support wrapping their corresponding kind of function.

We'll also want applying an AsyncContextDecorator to an async generator function to match whatever we decide for ContextDecorator on a sync generator.

Linked PRs

gh-125862: Context decorator support for generators and async functions #136212

This seems to work for me:

class ContextDecorator(object):
    "A base class or mixin that enables context managers to work as decorators."

    def _recreate_cm(self):
        """Return a recreated instance of self.

        Allows an otherwise one-shot context manager like
        _GeneratorContextManager to support use as
        a decorator via implicit recreation.

        This is a private interface just for _GeneratorContextManager.
        See issue #11647 for details.
        """
        return self

    def __call__(self, func):
        def inner(*args, **kwds):
            with self._recreate_cm():
                return func(*args, **kwds)

        def gen_inner(*args, **kwds):
            with self._recreate_cm():
                yield from func(*args, **kwds)

        async def async_inner(*args, **kwds):
            with self._recreate_cm():
                return await func(*args, **kwds)

        async def asyncgen_inner(*args, **kwds):
            with self._recreate_cm():
                async for value in func(*args, **kwds):
                    yield value

        wrapper = wraps(func)
        if isasyncgenfunction(func):
            return wrapper(asyncgen_inner)
        elif iscoroutinefunction(func):
            return wrapper(async_inner)
        elif isgeneratorfunction(func):
            return wrapper(gen_inner)
        else:
            return wrapper(inner)

Thoughts?

#114033 made a similar observation; my impression is that we can open a PR with @agronholm's proposal above and just fix this surprising behavior in Python 3.15 👍

The above changes may not work properly when context managers are being stacked and it is wrapping an already wrapped function, in which case the function type detection won't work as we'd like.

Gah, misclick and I don't have permission to reopen. Can a triager help out?

The above changes may not work properly when context managers are being stacked and it is wrapping an already wrapped function, in which case the function type detection won't work as we'd like.

Sounds like the first decorator should apply @inspect.markcoroutinefunction, no? I think stacking ContextDecorator will work correctly, since we're defining sync or async functions appropriately.

We could warn in most mistaken cases by replacing return func(*args, **kwds) with value = ...; if isawaitable(value): <issue warning> ; return value, but I could go either way on whether the performance overhead is worth it.

Sounds like the first decorator should apply @inspect.markcoroutinefunction, no? I think stacking ContextDecorator will work correctly, since we're defining sync or async functions appropriately.

But if our ContextDecorator is wrapping a decorator that wraps an async function, how can we possibly select the correct wrapper function?

istm that it's the responsibility of each decorator to return an appropriate function type, so if this doesn't work then I'd open the equivalent of this very issue against the decorator in question.

I've made the appropriate modifications to ContextDecorator and AsyncContextDecorator both, and their respective test suites still pass. Now I'm trying to come up with pathological cases to test against.

gh-125862: Keep ContextDecorator open across generator/coroutine exec…

fixed by #136212 - thanks, agronholm and gps!

Improve `contextlib.ContextDecorator` support for generators and async functions · Issue #125862 · python/cpython · GitHub

Repository navigation

Improve contextlib.ContextDecorator support for generators and async functions #125862

Linked PRs

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

Improve `contextlib.ContextDecorator` support for generators and async functions · Issue #125862 · python/cpython · GitHub

Repository navigation

Improve contextlib.ContextDecorator support for generators and async functions #125862

Description

Linked PRs

Activity

agronholm commented on May 13, 2025

Zac-HD commented on May 22, 2025

agronholm commented on May 24, 2025

Zac-HD commented on May 24, 2025

Zac-HD commented on May 24, 2025

agronholm commented on May 30, 2025

Zac-HD commented on May 31, 2025

agronholm commented on Jul 2, 2025

Zac-HD commented on Apr 28, 2026

Metadata

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

Issue actions