Uncancel tasks after timeout in python 3.11 (#363)

Author: tristan

CHANGES/362.feature

@@ -0,0 +1 @@
+Fix compatibility with asyncio.timeout() on python 3.11+.

async_timeout/__init__.py

@@ -12,6 +12,17 @@
     from typing_extensions import final
 
 
+if sys.version_info >= (3, 11):
+
+    def _uncancel_task(task: "asyncio.Task[object]") -> None:
+        task.uncancel()
+
+else:
+
+    def _uncancel_task(task: "asyncio.Task[object]") -> None:
+        pass
+
+
 __version__ = "4.0.2"
 
 
@@ -84,14 +95,15 @@ class Timeout:
     # The purpose is to time out as soon as possible
     # without waiting for the next await expression.
 
-    __slots__ = ("_deadline", "_loop", "_state", "_timeout_handler")
+    __slots__ = ("_deadline", "_loop", "_state", "_timeout_handler", "_task")
 
     def __init__(
         self, deadline: Optional[float], loop: asyncio.AbstractEventLoop
     ) -> None:
         self._loop = loop
         self._state = _State.INIT
 
+        self._task: Optional["asyncio.Task[object]"] = None
         self._timeout_handler = None  # type: Optional[asyncio.Handle]
         if deadline is None:
             self._deadline = None  # type: Optional[float]
@@ -147,6 +159,7 @@ def reject(self) -> None:
         self._reject()
 
     def _reject(self) -> None:
+        self._task = None
         if self._timeout_handler is not None:
             self._timeout_handler.cancel()
             self._timeout_handler = None
@@ -194,11 +207,11 @@ def _reschedule(self) -> None:
         if self._timeout_handler is not None:
             self._timeout_handler.cancel()
 
-        task = asyncio.current_task()
+        self._task = asyncio.current_task()
         if deadline <= now:
-            self._timeout_handler = self._loop.call_soon(self._on_timeout, task)
+            self._timeout_handler = self._loop.call_soon(self._on_timeout)
         else:
-            self._timeout_handler = self._loop.call_at(deadline, self._on_timeout, task)
+            self._timeout_handler = self._loop.call_at(deadline, self._on_timeout)
 
     def _do_enter(self) -> None:
         if self._state != _State.INIT:
@@ -208,15 +221,19 @@ def _do_enter(self) -> None:
 
     def _do_exit(self, exc_type: Optional[Type[BaseException]]) -> None:
         if exc_type is asyncio.CancelledError and self._state == _State.TIMEOUT:
+            assert self._task is not None
+            _uncancel_task(self._task)
             self._timeout_handler = None
+            self._task = None
             raise asyncio.TimeoutError
         # timeout has not expired
         self._state = _State.EXIT
         self._reject()
         return None
 
-    def _on_timeout(self, task: "asyncio.Task[None]") -> None:
-        task.cancel()
+    def _on_timeout(self) -> None:
+        assert self._task is not None
+        self._task.cancel()
         self._state = _State.TIMEOUT
         # drop the reference early
         self._timeout_handler = None

tests/test_timeout.py

@@ -1,4 +1,5 @@
 import asyncio
+import sys
 import time
 from functools import wraps
 from typing import Any, Callable, List, TypeVar
@@ -39,6 +40,10 @@ async def long_running_task() -> None:
             await long_running_task()
             assert t._loop is asyncio.get_event_loop()
     assert canceled_raised, "CancelledError was not raised"
+    if sys.version_info >= (3, 11):
+        task = asyncio.current_task()
+        assert task is not None
+        assert not task.cancelling()
 
 
 @pytest.mark.asyncio
@@ -158,6 +163,8 @@ async def outer() -> None:
     await task
     assert has_timeout
     assert not task.cancelled()
+    if sys.version_info >= (3, 11):
+        assert not task.cancelling()
     assert task.done()