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

Doc/library/contextlib.rst

@@ -467,12 +467,40 @@ Functions and classes provided:
       statements. If this is not the case, then the original construct with the
       explicit :keyword:`!with` statement inside the function should be used.
 
+   When the decorated callable is a generator function, coroutine function, or
+   asynchronous generator function, the returned wrapper is of the same kind
+   and keeps the context manager open for the lifetime of the iteration or
+   await rather than only for the call that creates the generator or coroutine
+   object.  Wrapped generators and asynchronous generators are explicitly
+   closed when iteration ends, as if by :func:`closing` or :func:`aclosing`.
+
+   .. note::
+      For asynchronous generators the wrapper re-yields each value with
+      ``async for``; values sent with :meth:`~agen.asend` and exceptions
+      thrown with :meth:`~agen.athrow` are not forwarded to the wrapped
+      generator.
+
    .. versionadded:: 3.2
 
+   .. versionchanged:: next
+      Decorating a generator function, coroutine function, or asynchronous
+      generator function now keeps the context manager open across iteration
+      or await.  Previously the context manager exited as soon as the
+      generator or coroutine object was created.
+
 
 .. class:: AsyncContextDecorator
 
-   Similar to :class:`ContextDecorator` but only for asynchronous functions.
+   Similar to :class:`ContextDecorator`, but the context manager is entered
+   and exited with :keyword:`async with`.  Decorate coroutine functions and
+   asynchronous generator functions with this class; the returned wrapper is
+   of the same kind.
+
+   .. note::
+      Synchronous functions and generators are accepted, but the wrapper is
+      always asynchronous, so the decorated callable must then be awaited or
+      iterated with ``async for``.  If that change of calling convention is
+      not intended, use :class:`ContextDecorator` instead.
 
    Example of ``AsyncContextDecorator``::
 
@@ -510,6 +538,13 @@ Functions and classes provided:
 
    .. versionadded:: 3.10
 
+   .. versionchanged:: next
+      Decorating an asynchronous generator function now keeps the context
+      manager open across iteration.  Previously the context manager exited
+      as soon as the generator object was created.  Synchronous functions
+      and synchronous generator functions are also now accepted, with an
+      asynchronous wrapper returned.
+
 
 .. class:: ExitStack()
 

Doc/whatsnew/3.15.rst

@@ -830,6 +830,15 @@ contextlib
   consistency with the :keyword:`with` and :keyword:`async with` statements.
   (Contributed by Serhiy Storchaka in :gh:`144386`.)
 
+* :class:`~contextlib.ContextDecorator` and
+  :class:`~contextlib.AsyncContextDecorator` (and therefore
+  :func:`~contextlib.contextmanager` and :func:`~contextlib.asynccontextmanager`
+  used as decorators) now detect generator functions, coroutine functions, and
+  asynchronous generator functions and keep the context manager open across
+  iteration or await.  Previously the context manager exited as soon as the
+  generator or coroutine object was created.
+  (Contributed by Alex Grönholm & Gregory P. Smith in :gh:`125862`.)
+
 
 dataclasses
 -----------

Lib/contextlib.py

@@ -1,10 +1,16 @@
 """Utilities for with-statement contexts.  See PEP 343."""
+
 import abc
 import os
 import sys
 import _collections_abc
 from collections import deque
 from functools import wraps
+lazy from inspect import (
+    isasyncgenfunction as _isasyncgenfunction,
+    iscoroutinefunction as _iscoroutinefunction,
+    isgeneratorfunction as _isgeneratorfunction,
+)
 from types import GenericAlias
 
 __all__ = ["asynccontextmanager", "contextmanager", "closing", "nullcontext",
@@ -79,11 +85,37 @@ def _recreate_cm(self):
         return self
 
     def __call__(self, func):
-        @wraps(func)
-        def inner(*args, **kwds):
-            with self._recreate_cm():
-                return func(*args, **kwds)
-        return inner
+        wrapper = wraps(func)
+        if _isasyncgenfunction(func):
+
+            async def asyncgen_inner(*args, **kwds):
+                with self._recreate_cm():
+                    async with aclosing(func(*args, **kwds)) as gen:
+                        async for value in gen:
+                            yield value
+
+            return wrapper(asyncgen_inner)
+        elif _iscoroutinefunction(func):
+
+            async def async_inner(*args, **kwds):
+                with self._recreate_cm():
+                    return await func(*args, **kwds)
+
+            return wrapper(async_inner)
+        elif _isgeneratorfunction(func):
+
+            def gen_inner(*args, **kwds):
+                with self._recreate_cm(), closing(func(*args, **kwds)) as gen:
+                    return (yield from gen)
+
+            return wrapper(gen_inner)
+        else:
+
+            def inner(*args, **kwds):
+                with self._recreate_cm():
+                    return func(*args, **kwds)
+
+            return wrapper(inner)
 
 
 class AsyncContextDecorator(object):
@@ -95,11 +127,41 @@ def _recreate_cm(self):
         return self
 
     def __call__(self, func):
-        @wraps(func)
-        async def inner(*args, **kwds):
-            async with self._recreate_cm():
-                return await func(*args, **kwds)
-        return inner
+        wrapper = wraps(func)
+        if _isasyncgenfunction(func):
+
+            async def asyncgen_inner(*args, **kwds):
+                async with (
+                    self._recreate_cm(),
+                    aclosing(func(*args, **kwds)) as gen
+                ):
+                    async for value in gen:
+                        yield value
+
+            return wrapper(asyncgen_inner)
+        elif _iscoroutinefunction(func):
+
+            async def async_inner(*args, **kwds):
+                async with self._recreate_cm():
+                    return await func(*args, **kwds)
+
+            return wrapper(async_inner)
+        elif _isgeneratorfunction(func):
+
+            async def gen_inner(*args, **kwds):
+                async with self._recreate_cm():
+                    with closing(func(*args, **kwds)) as gen:
+                        for value in gen:
+                            yield value
+
+            return wrapper(gen_inner)
+        else:
+
+            async def inner(*args, **kwds):
+                async with self._recreate_cm():
+                    return func(*args, **kwds)
+
+            return wrapper(inner)
 
 
 class _GeneratorContextManagerBase:

Lib/test/test_contextlib.py

@@ -680,6 +680,154 @@ def test(x):
         self.assertEqual(state, [1, 'something else', 999])
 
 
+    def test_contextmanager_decorate_generator_function(self):
+        @contextmanager
+        def woohoo(y):
+            state.append(y)
+            yield
+            state.append(999)
+
+        state = []
+        @woohoo(1)
+        def test(x):
+            self.assertEqual(state, [1])
+            state.append(x)
+            yield
+            state.append("second item")
+            return "result"
+
+        gen = test("something")
+        for _ in gen:
+            self.assertEqual(state, [1, "something"])
+        self.assertEqual(state, [1, "something", "second item", 999])
+
+        # The wrapped generator's return value is preserved.
+        state = []
+        gen = test("something")
+        with self.assertRaises(StopIteration) as cm:
+            while True:
+                next(gen)
+        self.assertEqual(cm.exception.value, "result")
+
+
+    def test_contextmanager_decorate_generator_function_exception(self):
+        @contextmanager
+        def woohoo():
+            state.append("enter")
+            try:
+                yield
+            finally:
+                state.append("exit")
+
+        state = []
+        @woohoo()
+        def test():
+            state.append("body")
+            yield
+            raise ZeroDivisionError
+
+        with self.assertRaises(ZeroDivisionError):
+            for _ in test():
+                pass
+        self.assertEqual(state, ["enter", "body", "exit"])
+
+
+    def test_contextmanager_decorate_generator_function_early_stop(self):
+        @contextmanager
+        def woohoo():
+            state.append("enter")
+            try:
+                yield
+            finally:
+                state.append("exit")
+
+        state = []
+        @woohoo()
+        def test():
+            try:
+                yield 1
+                yield 2
+            finally:
+                state.append("inner closed")
+
+        gen = test()
+        self.assertEqual(next(gen), 1)
+        gen.close()
+        # The inner generator is closed before the context manager exits.
+        self.assertEqual(state, ["enter", "inner closed", "exit"])
+
+
+    def test_contextmanager_decorate_generator_function_send_throw(self):
+        @contextmanager
+        def woohoo():
+            yield
+
+        @woohoo()
+        def test():
+            received = yield "first"
+            state.append(("received", received))
+            try:
+                yield "second"
+            except ValueError as exc:
+                state.append(("caught", type(exc)))
+                yield "after throw"
+
+        # .send() and .throw() are forwarded to the wrapped generator.
+        state = []
+        gen = test()
+        self.assertEqual(next(gen), "first")
+        self.assertEqual(gen.send("VALUE"), "second")
+        self.assertEqual(gen.throw(ValueError), "after throw")
+        gen.close()
+        self.assertEqual(
+            state, [("received", "VALUE"), ("caught", ValueError)]
+        )
+
+
+    def test_contextmanager_decorate_coroutine_function(self):
+        @contextmanager
+        def woohoo(y):
+            state.append(y)
+            yield
+            state.append(999)
+
+        state = []
+        @woohoo(1)
+        async def test(x):
+            self.assertEqual(state, [1])
+            state.append(x)
+
+        coro = test("something")
+        with self.assertRaises(StopIteration):
+            coro.send(None)
+
+        self.assertEqual(state, [1, "something", 999])
+
+
+    def test_contextmanager_decorate_asyncgen_function(self):
+        @contextmanager
+        def woohoo(y):
+            state.append(y)
+            yield
+            state.append(999)
+
+        state = []
+        @woohoo(1)
+        async def test(x):
+            self.assertEqual(state, [1])
+            state.append(x)
+            yield
+            state.append("second item")
+
+        agen = test("something")
+        with self.assertRaises(StopIteration):
+            agen.asend(None).send(None)
+        with self.assertRaises(StopAsyncIteration):
+            agen.asend(None).send(None)
+
+        self.assertEqual(state, [1, "something", "second item", 999])
+
+
 class TestBaseExitStack:
     exit_stack = None