Graceful contextmanager cleanup (#85)

* added tests for contextmanager handling GeneratorExit

* contextmanager handles closed child gracefully

Author: maxfischer2781

asyncstdlib/contextlib.py

@@ -90,7 +90,13 @@ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:
                 raise RuntimeError("generator did not stop after __aexit__")
         else:
             try:
-                await self.gen.athrow(exc_type, exc_val, exc_tb)
+                # We are being closed as part of (async) generator shutdown.
+                # Use `aclose` to have additional checks for the child to
+                # handle shutdown properly.
+                if exc_type is GeneratorExit:
+                    result = await self.gen.aclose()
+                else:
+                    result = await self.gen.athrow(exc_type, exc_val, exc_tb)
             except StopAsyncIteration as exc:
                 return exc is not exc_tb
             except RuntimeError as exc:
@@ -106,6 +112,16 @@ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:
                     raise
                 return False
             else:
+                # During shutdown, the child generator might be cleaned up early.
+                # In this case,
+                # - the child will return nothing/None,
+                # - we get cleaned up via GeneratorExit as well,
+                # and we should go on with our own cleanup.
+                # This might happen if the child mishandles GeneratorExit as well,
+                # but is the closest we can get to checking the situation.
+                # See https://github.com/maxfischer2781/asyncstdlib/issues/84
+                if exc_type is GeneratorExit and result is None:
+                    return False
                 raise RuntimeError("generator did not stop after throw() in __aexit__")
 
 

unittests/test_contextlib.py

@@ -34,6 +34,8 @@ async def valid(value):
 
 @sync
 async def test_contextmanager_no_yield():
+    """Test that it is an error for a context to not yield"""
+
     @a.contextmanager
     async def no_yield():
         if False:
@@ -46,6 +48,8 @@ async def no_yield():
 
 @sync
 async def test_contextmanager_no_stop():
+    """Test that it is an error for a context to yield again after stopping"""
+
     @a.contextmanager
     async def no_stop():
         yield
@@ -69,6 +73,8 @@ async def supress_no_stop():
 
 @sync
 async def test_contextmanager_raise_asyncstop():
+    """Test that StopAsyncIteration may propagate out of a context block"""
+
     @a.contextmanager
     async def no_raise():
         yield
@@ -113,6 +119,8 @@ async def replace():
 
 @sync
 async def test_contextmanager_raise_same():
+    """Test that outer exceptions do not shadow inner/newer ones"""
+
     @a.contextmanager
     async def reraise():
         try:
@@ -136,6 +144,65 @@ async def recreate():
             raise KeyError("outside")
 
 
+@sync
+async def test_contextmanager_raise_generatorexit():
+    """Test that shutdown via GeneratorExit is propagated"""
+
+    @a.contextmanager
+    async def no_op():
+        yield
+
+    with pytest.raises(GeneratorExit):
+        async with no_op():
+            raise GeneratorExit("used to tear down coroutines")
+
+    # during shutdown, generators may be killed in arbitrary order
+    # make sure we do not suppress GeneratorExit
+    with pytest.raises(GeneratorExit, match="inner"):
+        context = no_op()
+        async with context:
+            # simulate cleanup closing the child early
+            await context.gen.aclose()
+            raise GeneratorExit("inner")
+
+
+@sync
+async def test_contextmanager_no_suppress_generatorexit():
+    """Test that GeneratorExit is not suppressed"""
+
+    @a.contextmanager
+    async def no_op():
+        yield
+
+    exc = GeneratorExit("GE should not be replaced normally")
+    with pytest.raises(type(exc)) as exc_info:
+        async with no_op():
+            raise exc
+    assert exc_info.value is exc
+
+    @a.contextmanager
+    async def exit_ge():
+        try:
+            yield
+        except GeneratorExit:
+            pass
+
+    with pytest.raises(GeneratorExit):
+        async with exit_ge():
+            raise GeneratorExit("Resume teardown if child exited")
+
+    @a.contextmanager
+    async def ignore_ge():
+        try:
+            yield
+        except GeneratorExit:
+            yield
+
+    with pytest.raises(RuntimeError):
+        async with ignore_ge():
+            raise GeneratorExit("Warn if child does not exit")
+
+
 @sync
 async def test_nullcontext():
     async with a.nullcontext(1337) as value: