Restore timeout cancellation state

Author: 8r1ef

asgiref/timeout.py

@@ -7,6 +7,7 @@
 
 
 import asyncio
+import sys
 import warnings
 from types import TracebackType
 from typing import Any  # noqa
@@ -43,6 +44,7 @@ def __init__(
             )
         self._loop = loop
         self._task = None  # type: Optional[asyncio.Task[Any]]
+        self._cancelling = 0
         self._cancelled = False
         self._cancel_handler = None  # type: Optional[asyncio.Handle]
         self._cancel_at = None  # type: Optional[float]
@@ -92,6 +94,8 @@ def _do_enter(self) -> "timeout":
             raise RuntimeError(
                 "Timeout context manager should be used " "inside a task"
             )
+        if sys.version_info >= (3, 11):
+            self._cancelling = self._task.cancelling()
 
         if self._timeout <= 0:
             self._loop.call_soon(self._cancel_task)
@@ -102,7 +106,11 @@ def _do_enter(self) -> "timeout":
         return self
 
     def _do_exit(self, exc_type: Type[BaseException]) -> None:
-        if exc_type is asyncio.CancelledError and self._cancelled:
+        should_raise_timeout = self._cancelled
+        if self._cancelled and sys.version_info >= (3, 11):
+            assert self._task is not None
+            should_raise_timeout = self._task.uncancel() <= self._cancelling
+        if exc_type is asyncio.CancelledError and should_raise_timeout:
             self._cancel_handler = None
             self._task = None
             raise asyncio.TimeoutError

tests/test_timeout.py

@@ -0,0 +1,61 @@
+import asyncio
+import sys
+
+import pytest
+
+from asgiref.timeout import timeout
+
+
+@pytest.mark.asyncio
+async def test_timeout_restores_cancellation_count() -> None:
+    if sys.version_info < (3, 11):
+        pytest.skip("Task cancellation counts were added in Python 3.11")
+
+    task = asyncio.current_task()
+    assert task is not None
+    cancelling = task.cancelling()
+
+    with pytest.raises(asyncio.TimeoutError):
+        async with timeout(0):
+            await asyncio.sleep(0)
+
+    assert task.cancelling() == cancelling
+
+
+@pytest.mark.asyncio
+async def test_timeout_restores_count_when_cancellation_is_suppressed() -> None:
+    if sys.version_info < (3, 11):
+        pytest.skip("Task cancellation counts were added in Python 3.11")
+
+    task = asyncio.current_task()
+    assert task is not None
+    cancelling = task.cancelling()
+
+    async with timeout(0):
+        try:
+            await asyncio.sleep(0)
+        except asyncio.CancelledError:
+            pass
+
+    assert task.cancelling() == cancelling
+
+
+@pytest.mark.asyncio
+async def test_timeout_preserves_external_cancellation() -> None:
+    if sys.version_info < (3, 11):
+        pytest.skip("Task cancellation counts were added in Python 3.11")
+
+    task = asyncio.current_task()
+    assert task is not None
+    cancelling = task.cancelling()
+
+    try:
+        with pytest.raises(asyncio.CancelledError):
+            async with timeout(0):
+                asyncio.get_running_loop().call_soon(task.cancel)
+                await asyncio.sleep(0)
+
+        assert task.cancelling() == cancelling + 1
+    finally:
+        while task.cancelling() > cancelling:
+            task.uncancel()