gh-102780: Fix uncancel() call in asyncio timeouts (#102815)

Also use `raise TimeOut from <CancelledError instance>` so that the CancelledError is set
in the `__cause__` field rather than in the `__context__` field.

Co-authored-by: Guido van Rossum <gvanrossum@gmail.com>
Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>

2 files changed, 34 insertions, 3 deletions

diff --git a/Lib/asyncio/timeouts.py b/Lib/asyncio/timeouts.py
index d07b291..029c468 100644
--- a/Lib/asyncio/timeouts.py
+++ b/Lib/asyncio/timeouts.py
@@ -84,6 +84,7 @@ class Timeout:
     async def __aenter__(self) -> "Timeout":
         self._state = _State.ENTERED
         self._task = tasks.current_task()
+        self._cancelling = self._task.cancelling()
         if self._task is None:
             raise RuntimeError("Timeout should be used inside a task")
         self.reschedule(self._when)
@@ -104,10 +105,10 @@ class Timeout:
         if self._state is _State.EXPIRING:
             self._state = _State.EXPIRED
 
-            if self._task.uncancel() == 0 and exc_type is exceptions.CancelledError:
-                # Since there are no outstanding cancel requests, we're
+            if self._task.uncancel() <= self._cancelling and exc_type is exceptions.CancelledError:
+                # Since there are no new cancel requests, we're
                 # handling this.
-                raise TimeoutError
+                raise TimeoutError from exc_val
         elif self._state is _State.ENTERED:
             self._state = _State.EXITED
 
diff --git a/Lib/test/test_asyncio/test_timeouts.py b/Lib/test/test_asyncio/test_timeouts.py
index b9bac6f..8b6b9a1 100644
--- a/Lib/test/test_asyncio/test_timeouts.py
+++ b/Lib/test/test_asyncio/test_timeouts.py
@@ -247,6 +247,36 @@ class TimeoutTests(unittest.IsolatedAsyncioTestCase):
                         async with asyncio.timeout(0.01):
                             await asyncio.sleep(10)
 
+    async def test_timeout_after_cancellation(self):
+        try:
+            asyncio.current_task().cancel()
+            await asyncio.sleep(1)  # work which will be cancelled
+        except asyncio.CancelledError:
+            pass
+        finally:
+            with self.assertRaises(TimeoutError):
+                async with asyncio.timeout(0.0):
+                    await asyncio.sleep(1)  # some cleanup
+
+    async def test_cancel_in_timeout_after_cancellation(self):
+        try:
+            asyncio.current_task().cancel()
+            await asyncio.sleep(1)  # work which will be cancelled
+        except asyncio.CancelledError:
+            pass
+        finally:
+            with self.assertRaises(asyncio.CancelledError):
+                async with asyncio.timeout(1.0):
+                    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)
+
 
 if __name__ == '__main__':
     unittest.main()