Fix regression: (#258)

* Fix regression:
1. Don't raise TimeoutError from timeout object that doesn't enter into async context manager
2. Use call_soon() for raising TimeoutError if deadline is reached on entering into async context manager

CHANGES.rst

@@ -1,6 +1,19 @@
 CHANGES
 =======
 
+4.0.1 (2121-11-xx)
+------------------
+
+- Fix regression:
+
+  1. Don't raise TimeoutError from timeout object that doesn't enter into async context
+     manager
+
+  2. Use call_soon() for raising TimeoutError if deadline is reached on entering into
+     async context manager
+
+  (#258)
+
 4.0.0 (2021-11-01)
 ------------------
 

async_timeout/__init__.py

@@ -80,17 +80,14 @@ class Timeout:
     # The purpose is to time out as sson as possible
     # without waiting for the next await expression.
 
-    __slots__ = ("_deadline", "_loop", "_state", "_task", "_timeout_handler")
+    __slots__ = ("_deadline", "_loop", "_state", "_timeout_handler")
 
     def __init__(
         self, deadline: Optional[float], loop: asyncio.AbstractEventLoop
     ) -> None:
         self._loop = loop
         self._state = _State.INIT
 
-        task = _current_task(self._loop)
-        self._task = task
-
         self._timeout_handler = None  # type: Optional[asyncio.Handle]
         if deadline is None:
             self._deadline = None  # type: Optional[float]
@@ -180,22 +177,30 @@ def update(self, deadline: float) -> None:
         if self._timeout_handler is not None:
             self._timeout_handler.cancel()
         self._deadline = deadline
+        if self._state != _State.INIT:
+            self._reschedule()
+
+    def _reschedule(self) -> None:
+        assert self._state == _State.ENTER
+        deadline = self._deadline
+        if deadline is None:
+            return
+
         now = self._loop.time()
+        if self._timeout_handler is not None:
+            self._timeout_handler.cancel()
+
+        task = _current_task(self._loop)
         if deadline <= now:
-            self._timeout_handler = None
-            if self._state == _State.INIT:
-                raise asyncio.TimeoutError
-            else:
-                # state is ENTER
-                raise asyncio.CancelledError
-        self._timeout_handler = self._loop.call_at(
-            deadline, self._on_timeout, self._task
-        )
+            self._timeout_handler = self._loop.call_soon(self._on_timeout, task)
+        else:
+            self._timeout_handler = self._loop.call_at(deadline, self._on_timeout, task)
 
     def _do_enter(self) -> None:
         if self._state != _State.INIT:
             raise RuntimeError(f"invalid state {self._state.value}")
         self._state = _State.ENTER
+        self._reschedule()
 
     def _do_exit(self, exc_type: Type[BaseException]) -> None:
         if exc_type is asyncio.CancelledError and self._state == _State.TIMEOUT:
@@ -209,6 +214,8 @@ def _do_exit(self, exc_type: Type[BaseException]) -> None:
     def _on_timeout(self, task: "asyncio.Task[None]") -> None:
         task.cancel()
         self._state = _State.TIMEOUT
+        # drop the reference early
+        self._timeout_handler = None
 
 
 if sys.version_info >= (3, 7):

tests/test_timeout.py

@@ -82,7 +82,8 @@ def test_timeout_no_loop() -> None:
 @pytest.mark.asyncio
 async def test_timeout_zero() -> None:
     with pytest.raises(asyncio.TimeoutError):
-        timeout(0)
+        async with timeout(0):
+            await asyncio.sleep(10)
 
 
 @pytest.mark.asyncio
@@ -307,10 +308,11 @@ async def test_shift_nonscheduled() -> None:
 
 
 @pytest.mark.asyncio
-async def test_shift_by_negative_expired() -> None:
+async def test_shift_negative_expired() -> None:
     async with timeout(1) as cm:
         with pytest.raises(asyncio.CancelledError):
             cm.shift(-1)
+            await asyncio.sleep(10)
 
 
 @pytest.mark.asyncio