gh-95882 fix traceback of exceptions propagated from inside a contextlib.asynccontextmanager

Author: graingert

Lib/contextlib.py

@@ -228,6 +228,7 @@ async def __aexit__(self, typ, value, traceback):
             except RuntimeError as exc:
                 # Don't re-raise the passed in exception. (issue27122)
                 if exc is value:
+                    exc.__traceback__ = traceback
                     return False
                 # Avoid suppressing if a Stop(Async)Iteration exception
                 # was passed to athrow() and later wrapped into a RuntimeError
@@ -239,6 +240,7 @@ async def __aexit__(self, typ, value, traceback):
                     isinstance(value, (StopIteration, StopAsyncIteration))
                     and exc.__cause__ is value
                 ):
+                    exc.__traceback__ = traceback
                     return False
                 raise
             except BaseException as exc:
@@ -250,6 +252,7 @@ async def __aexit__(self, typ, value, traceback):
                 # and the __exit__() protocol.
                 if exc is not value:
                     raise
+                exc.__traceback__ = traceback
                 return False
             raise RuntimeError("generator didn't stop after athrow()")
 

Lib/test/test_contextlib_async.py

@@ -125,6 +125,33 @@ async def woohoo():
                 raise ZeroDivisionError()
         self.assertEqual(state, [1, 42, 999])
 
+    @_async_test
+    async def test_contextmanager_traceback(self):
+        @asynccontextmanager
+        async def f():
+            yield
+
+        try:
+            async with f():
+                1/0
+        except ZeroDivisionError as e:
+            frames = traceback.extract_tb(e.__traceback__)
+
+        self.assertEqual(len(frames), 1)
+        self.assertEqual(frames[0].name, 'test_contextmanager_traceback')
+        self.assertEqual(frames[0].line, '1/0')
+
+        # Repeat with RuntimeError (which goes through a different code path)
+        try:
+            async with f():
+                raise NotImplementedError(42)
+        except NotImplementedError as e:
+            frames = traceback.extract_tb(e.__traceback__)
+
+        self.assertEqual(len(frames), 1)
+        self.assertEqual(frames[0].name, 'test_contextmanager_traceback')
+        self.assertEqual(frames[0].line, 'raise NotImplementedError(42)')
+
     @_async_test
     async def test_contextmanager_no_reraise(self):
         @asynccontextmanager