asyncio: #10 shutdown_asyncgens

Ensure any async generators that are left running past
run_until_complete (essentially only if we get a KeyboardInterrupt that
stops the run_until_complete) then those are finalized properly

Note I don't use loop.shutdown_asyncgens sothat I can make the generator
handle an asyncio.CancelledError rather than a GeneratorExit

Author: delfick

README.rst

@@ -25,6 +25,7 @@ Changelog
     * Fix bug where it was possible for an async generator fixture to
       be cleaned up even if it was never started.
     * This library is now 3.7+ only
+    * Added an equivalent ``shutdown_asyncgen`` to the OverrideLoop helper
 
 0.5.4 - 26 January 2021
     * Added a ``--default-async-timeout`` option from the CLI. With many thanks

alt_pytest_asyncio/plugin.py

@@ -199,11 +199,46 @@ def __exit__(self, exc_typ, exc, tb):
         try:
             if getattr(self, "loop", None):
                 cancel_all_tasks(self.loop, ignore_errors_from_tasks=self.tasks)
+                self.loop.run_until_complete(self.shutdown_asyncgens())
                 self.loop.close()
         finally:
             if hasattr(self, "_original_loop"):
                 asyncio.set_event_loop(self._original_loop)
 
+    async def shutdown_asyncgens(self):
+        """
+        A version of loop.shutdown_asyncgens that tries to cancel the generators
+        before closing them.
+        """
+        if not len(self.loop._asyncgens):
+            return
+
+        closing_agens = list(self.loop._asyncgens)
+        self.loop._asyncgens.clear()
+
+        # I would do an asyncio.tasks.gather but it would appear that just causes
+        # the asyncio loop to think it's shutdown, so I have to do them one at a time
+        for ag in closing_agens:
+            try:
+                try:
+                    try:
+                        await ag.athrow(asyncio.CancelledError())
+                    except StopAsyncIteration:
+                        pass
+                finally:
+                    await ag.aclose()
+            except asyncio.CancelledError:
+                pass
+            except:
+                exc = sys.exc_info()[1]
+                self.loop.call_exception_handler(
+                    {
+                        "message": "an error occurred during closing of asynchronous generator",
+                        "exception": exc,
+                        "asyncgen": ag,
+                    }
+                )
+
     def run_until_complete(self, coro):
         if not hasattr(self, "loop"):
             raise Exception(

pylama.ini

@@ -2,7 +2,7 @@
 skip = */setup.py
 
 [pylama:tests/*]
-ignore = E225,E202,E211,E231,E226,W292,W291,E251,E122,E501,E701,E227,E305,E128,W391
+ignore = E225,E202,E211,E231,E226,W292,W291,E251,E122,E501,E701,E227,E305,E128,W391,C901
 
 [pylama:pyflakes]
 builtins = _

tests/test_override_loop.py

@@ -137,6 +137,78 @@ async def blah():
         get_event_loop().run_until_complete(info["coro"])
 
 
+it "can shutdown async gens":
+    info1 = []
+    info2 = []
+    info3 = []
+
+    original = asyncio.get_event_loop()
+
+    async def my_generator(info):
+        try:
+            info.append(1)
+            yield
+            info.append(2)
+            yield
+            info.append(3)
+        except asyncio.CancelledError:
+            info.append("cancelled")
+            raise
+        finally:
+            info.append(("done", __import__("sys").exc_info()[0]))
+
+    # Test that the outside loop isn't affected by the inside loop
+    outside_gen = my_generator(info1)
+
+    async def outside1():
+        await outside_gen.__anext__()
+        await outside_gen.__anext__()
+
+    original.run_until_complete(outside1())
+    assert info1 == [1, 2]
+
+    # The way python asyncio works
+    # Means that by defining this outside our OverrideLoop
+    # The weakref held against it in the _asyncgens set on the loop
+    # Will remain so that our shutdown_asyncgens function may work
+    ag = my_generator(info2)
+
+    with OverrideLoop(new_loop=True) as custom_loop:
+        assert info2 == []
+        assert info3 == []
+
+        async def doit():
+            ag2 = my_generator(info3)
+            assert set(asyncio.get_event_loop()._asyncgens) == set()
+            await ag2.__anext__()
+            assert set(asyncio.get_event_loop()._asyncgens) == set([ag2])
+            await ag.__anext__()
+            assert set(asyncio.get_event_loop()._asyncgens) == set([ag2, ag])
+            await ag.__anext__()
+            assert info3 == [1]
+
+        custom_loop.run_until_complete(doit())
+        assert list(custom_loop.loop._asyncgens) == [ag]
+        assert info3 == [1]
+        assert info2 == [1, 2]
+
+    assert asyncio.get_event_loop() is original
+    assert not original.is_closed()
+
+    assert info3 == [1, "cancelled", ("done", asyncio.CancelledError)]
+    assert info2 == [1, 2, "cancelled", ("done", asyncio.CancelledError)]
+    assert info1 == [1, 2]
+
+    async def outside2():
+        try:
+            await outside_gen.__anext__()
+        except StopAsyncIteration:
+            pass
+
+    # Test that the outside loop isn't affected by the inside loop
+    original.run_until_complete(outside2())
+    assert info1 == [1, 2, 3, ("done", None)]
+
 describe "testing autouse":
 
     @pytest.fixture(autouse=True)