gh-125862: Improve context decorator support for generators and async functions

Author: agronholm

Lib/contextlib.py

@@ -1,4 +1,7 @@
 """Utilities for with-statement contexts.  See PEP 343."""
+from inspect import isasyncgenfunction, iscoroutinefunction, \
+    isgeneratorfunction
+
 import abc
 import os
 import sys
@@ -79,11 +82,32 @@ 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
+
+        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)
 
 
 class AsyncContextDecorator(object):
@@ -95,11 +119,33 @@ def _recreate_cm(self):
         return self
 
     def __call__(self, func):
-        @wraps(func)
         async def inner(*args, **kwds):
+            async with self._recreate_cm():
+                return func(*args, **kwds)
+
+        async def gen_inner(*args, **kwds):
+            async with self._recreate_cm():
+                for value in func(*args, **kwds):
+                    yield value
+
+        async def async_inner(*args, **kwds):
             async with self._recreate_cm():
                 return await func(*args, **kwds)
-        return inner
+
+        async def asyncgen_inner(*args, **kwds):
+            async 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)
 
 
 class _GeneratorContextManagerBase:

Lib/test/test_contextlib.py

@@ -1,5 +1,5 @@
 """Unit tests for contextlib.py, and other context managers."""
-
+import asyncio
 import io
 import os
 import sys
@@ -680,6 +680,66 @@ 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")
+
+        for _ in test("something"):
+            self.assertEqual(state, [1, "something"])
+        self.assertEqual(state, [1, "something", "second item", 999])
+
+
+    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)
+
+        asyncio.run(test('something'))
+        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")
+
+        async def run_test():
+            async for _ in test("something"):
+                self.assertEqual(state, [1, "something"])
+
+        asyncio.run(run_test())
+        self.assertEqual(state, [1, 'something', "second item", 999])
+
+
 class TestBaseExitStack:
     exit_stack = None
 

Lib/test/test_contextlib_async.py

@@ -402,6 +402,63 @@ async def test():
         await test()
         self.assertFalse(entered)
 
+    @_async_test
+    async def test_decorator_decorate_sync_function(self):
+        @asynccontextmanager
+        async def context():
+            state.append(1)
+            yield
+            state.append(999)
+
+        state = []
+        @context()
+        def test(x):
+            self.assertEqual(state, [1])
+            state.append(x)
+
+        await test("something")
+        self.assertEqual(state, [1, "something", 999])
+
+    @_async_test
+    async def test_decorator_decorate_generator_function(self):
+        @asynccontextmanager
+        async def context():
+            state.append(1)
+            yield
+            state.append(999)
+
+        state = []
+        @context()
+        def test(x):
+            self.assertEqual(state, [1])
+            state.append(x)
+            yield
+            state.append("second item")
+
+        async for _ in test("something"):
+            self.assertEqual(state, [1, "something"])
+        self.assertEqual(state, [1, "something", "second item", 999])
+
+    @_async_test
+    async def test_decorator_decorate_asyncgen_function(self):
+        @asynccontextmanager
+        async def context():
+            state.append(1)
+            yield
+            state.append(999)
+
+        state = []
+        @context()
+        async def test(x):
+            self.assertEqual(state, [1])
+            state.append(x)
+            yield
+            state.append("second item")
+
+        async for _ in test("something"):
+            self.assertEqual(state, [1, "something"])
+        self.assertEqual(state, [1, "something", "second item", 999])
+
     @_async_test
     async def test_decorator_with_exception(self):
         entered = False