gh-96037: Always insert TimeoutError when exit an expired asyncio.timeout() block

Lib/asyncio/timeouts.py

@@ -109,10 +109,16 @@ async def __aexit__(
         if self._state is _State.EXPIRING:
             self._state = _State.EXPIRED
 
-            if self._task.uncancel() <= self._cancelling and exc_type is exceptions.CancelledError:
+            if self._task.uncancel() <= self._cancelling:
                 # Since there are no new cancel requests, we're
                 # handling this.
-                raise TimeoutError from exc_val
+                if exc_type is exceptions.CancelledError:
+                    raise TimeoutError from exc_val
+                elif exc_val is not None:
+                    self._insert_timeout_error(exc_val)
+                    if isinstance(exc_val, ExceptionGroup):
+                        for exc in exc_val.exceptions:
+                            self._insert_timeout_error(exc)
         elif self._state is _State.ENTERED:
             self._state = _State.EXITED
 
@@ -125,6 +131,16 @@ def _on_timeout(self) -> None:
         # drop the reference early
         self._timeout_handler = None
 
+    @staticmethod
+    def _insert_timeout_error(exc_val: BaseException) -> None:
+        while exc_val.__context__ is not None:
+            if type(exc_val.__context__) is exceptions.CancelledError:
+                te = TimeoutError()
+                te.__context__ = te.__cause__ = exc_val.__context__
+                exc_val.__context__ = te
+                break
+            exc_val = exc_val.__context__
+
 
 def timeout(delay: Optional[float]) -> Timeout:
     """Timeout async context manager.

Lib/test/test_asyncio/test_timeouts.py

@@ -116,15 +116,68 @@ async def test_foreign_exception_passed(self):
                 raise KeyError
         self.assertFalse(cm.expired())
 
+    async def test_timeout_exception_context(self):
+        with self.assertRaises(TimeoutError) as cm:
+            async with asyncio.timeout(0.01):
+                try:
+                    1/0
+                finally:
+                    await asyncio.sleep(1)
+        e = cm.exception
+        # Expect TimeoutError caused by CancelledError raised during handling
+        # of ZeroDivisionError.
+        e2 = e.__cause__
+        self.assertIsInstance(e2, asyncio.CancelledError)
+        self.assertIs(e.__context__, e2)
+        self.assertIsNone(e2.__cause__)
+        self.assertIsInstance(e2.__context__, ZeroDivisionError)
+
     async def test_foreign_exception_on_timeout(self):
         async def crash():
             try:
                 await asyncio.sleep(1)
             finally:
                 1/0
-        with self.assertRaises(ZeroDivisionError):
+        with self.assertRaises(ZeroDivisionError) as cm:
             async with asyncio.timeout(0.01):
                 await crash()
+        e = cm.exception
+        # Expect ZeroDivisionError raised during handling of TimeoutError
+        # caused by CancelledError.
+        self.assertIsNone(e.__cause__)
+        e2 = e.__context__
+        self.assertIsInstance(e2, TimeoutError)
+        e3 = e2.__cause__
+        self.assertIsInstance(e3, asyncio.CancelledError)
+        self.assertIs(e2.__context__, e3)
+
+    async def test_foreign_exception_on_timeout_2(self):
+        with self.assertRaises(ZeroDivisionError) as cm:
+            async with asyncio.timeout(0.01):
+                try:
+                    try:
+                        raise ValueError
+                    finally:
+                        await asyncio.sleep(1)
+                finally:
+                    try:
+                        raise KeyError
+                    finally:
+                        1/0
+        e = cm.exception
+        # Expect ZeroDivisionError raised during handling of KeyError
+        # raised during handling of TimeoutError caused by CancelledError.
+        self.assertIsNone(e.__cause__)
+        e2 = e.__context__
+        self.assertIsInstance(e2, KeyError)
+        self.assertIsNone(e2.__cause__)
+        e3 = e2.__context__
+        self.assertIsInstance(e3, TimeoutError)
+        e4 = e3.__cause__
+        self.assertIsInstance(e4, asyncio.CancelledError)
+        self.assertIsNone(e4.__cause__)
+        self.assertIsInstance(e4.__context__, ValueError)
+        self.assertIs(e3.__context__, e4)
 
     async def test_foreign_cancel_doesnt_timeout_if_not_expired(self):
         with self.assertRaises(asyncio.CancelledError):
@@ -219,14 +272,30 @@ async def test_repr_disabled(self):
             self.assertEqual(repr(cm), r"<Timeout [active] when=None>")
 
     async def test_nested_timeout_in_finally(self):
-        with self.assertRaises(TimeoutError):
+        with self.assertRaises(TimeoutError) as cm1:
             async with asyncio.timeout(0.01):
                 try:
                     await asyncio.sleep(1)
                 finally:
-                    with self.assertRaises(TimeoutError):
+                    with self.assertRaises(TimeoutError) as cm2:
                         async with asyncio.timeout(0.01):
                             await asyncio.sleep(10)
+        e1 = cm1.exception
+        # Expect TimeoutError caused by CancelledError.
+        e12 = e1.__cause__
+        self.assertIsInstance(e12, asyncio.CancelledError)
+        self.assertIsNone(e12.__cause__)
+        self.assertIsNone(e12.__context__)
+        self.assertIs(e1.__context__, e12)
+        e2 = cm2.exception
+        # Expect TimeoutError caused by CancelledError raised during
+        # handling of other CancelledError (which is the same as in
+        # the above chain).
+        e22 = e2.__cause__
+        self.assertIsInstance(e22, asyncio.CancelledError)
+        self.assertIsNone(e22.__cause__)
+        self.assertIs(e22.__context__, e12)
+        self.assertIs(e2.__context__, e22)
 
     async def test_timeout_after_cancellation(self):
         try:
@@ -235,7 +304,7 @@ async def test_timeout_after_cancellation(self):
         except asyncio.CancelledError:
             pass
         finally:
-            with self.assertRaises(TimeoutError):
+            with self.assertRaises(TimeoutError) as cm:
                 async with asyncio.timeout(0.0):
                     await asyncio.sleep(1)  # some cleanup
 
@@ -251,13 +320,6 @@ async def test_cancel_in_timeout_after_cancellation(self):
                     asyncio.current_task().cancel()
                     await asyncio.sleep(2)  # some cleanup
 
-    async def test_timeout_exception_cause (self):
-        with self.assertRaises(asyncio.TimeoutError) as exc:
-            async with asyncio.timeout(0):
-                await asyncio.sleep(1)
-        cause = exc.exception.__cause__
-        assert isinstance(cause, asyncio.CancelledError)
-
     async def test_timeout_already_entered(self):
         async with asyncio.timeout(0.01) as cm:
             with self.assertRaisesRegex(RuntimeError, "has already been entered"):
@@ -303,6 +365,47 @@ async def test_timeout_without_task(self):
         with self.assertRaisesRegex(RuntimeError, "has not been entered"):
             cm.reschedule(0.02)
 
+    async def test_timeout_taskgroup(self):
+        async def task():
+            try:
+                await asyncio.sleep(2)  # Will be interrupted after 0.01 second
+            finally:
+                1/0  # Crash in cleanup
+
+        with self.assertRaises(ExceptionGroup) as cm:
+            async with asyncio.timeout(0.01):
+                async with asyncio.TaskGroup() as tg:
+                    tg.create_task(task())
+                    try:
+                        raise ValueError
+                    finally:
+                        await asyncio.sleep(1)
+        eg = cm.exception
+        # Expect ExceptionGroup raised during handling of TimeoutError caused
+        # by CancelledError raised during handling of ValueError.
+        self.assertIsNone(eg.__cause__)
+        e_1 = eg.__context__
+        self.assertIsInstance(e_1, TimeoutError)
+        e_2 = e_1.__cause__
+        self.assertIsInstance(e_2, asyncio.CancelledError)
+        self.assertIsNone(e_2.__cause__)
+        self.assertIsInstance(e_2.__context__, ValueError)
+        self.assertIs(e_1.__context__, e_2)
+
+        self.assertEqual(len(eg.exceptions), 1, eg)
+        e1 = eg.exceptions[0]
+        # Expect ZeroDivisionError raised during handling of TimeoutError
+        # caused by CancelledError (it is a different CancelledError).
+        self.assertIsInstance(e1, ZeroDivisionError)
+        self.assertIsNone(e1.__cause__)
+        e2 = e1.__context__
+        self.assertIsInstance(e2, TimeoutError)
+        e3 = e2.__cause__
+        self.assertIsInstance(e3, asyncio.CancelledError)
+        self.assertIsNone(e3.__context__)
+        self.assertIsNone(e3.__cause__)
+        self.assertIs(e2.__context__, e3)
+
 
 if __name__ == '__main__':
     unittest.main()

Misc/NEWS.d/next/Library/2024-01-08-19-38-42.gh-issue-96037.Yr2Y1C.rst

@@ -0,0 +1,2 @@
+Insert :exc:`TimeoutError` in the context of the exception that was raised
+during exiting an expired :func:`asyncio.timeout` block.