Type completeness improvement (timeouts, CancelScope, and more) (#2671)

* Add --full-diagnostics-file to check_type_completeness.py for use when developing

* _timeouts.py, trio/_core, trio/_tests, open_nursery, move_on_at, move_on_after and CancelScope now type complete

trio/_core/_exceptions.py

@@ -59,7 +59,7 @@ class Cancelled(BaseException, metaclass=NoPublicConstructor):
 
     """
 
-    def __str__(self):
+    def __str__(self) -> str:
         return "Cancelled"
 
 

trio/_core/_run.py

@@ -17,6 +17,9 @@
 from math import inf
 from time import perf_counter
 from typing import TYPE_CHECKING, Any, NoReturn, TypeVar
+from types import TracebackType
+from collections.abc import Iterator
+from contextlib import AbstractAsyncContextManager
 
 import attr
 from outcome import Error, Outcome, Value, capture
@@ -475,15 +478,15 @@ class CancelScope(metaclass=Final):
     has been entered yet, and changes take immediate effect.
     """
 
-    _cancel_status = attr.ib(default=None, init=False)
-    _has_been_entered = attr.ib(default=False, init=False)
-    _registered_deadline = attr.ib(default=inf, init=False)
-    _cancel_called = attr.ib(default=False, init=False)
-    cancelled_caught = attr.ib(default=False, init=False)
+    _cancel_status: CancelStatus | None = attr.ib(default=None, init=False)
+    _has_been_entered: bool = attr.ib(default=False, init=False)
+    _registered_deadline: float = attr.ib(default=inf, init=False)
+    _cancel_called: bool = attr.ib(default=False, init=False)
+    cancelled_caught: bool = attr.ib(default=False, init=False)
 
     # Constructor arguments:
-    _deadline = attr.ib(default=inf, kw_only=True)
-    _shield = attr.ib(default=False, kw_only=True)
+    _deadline: float = attr.ib(default=inf, kw_only=True)
+    _shield: bool = attr.ib(default=False, kw_only=True)
 
     @enable_ki_protection
     def __enter__(self):
@@ -573,7 +576,12 @@ def _close(self, exc):
             self._cancel_status = None
         return exc
 
-    def __exit__(self, etype, exc, tb):
+    def __exit__(
+        self,
+        etype: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> bool:
         # NB: NurseryManager calls _close() directly rather than __exit__(),
         # so __exit__() must be just _close() plus this logic for adapting
         # the exception-filtering result to the context manager API.
@@ -607,7 +615,7 @@ def __exit__(self, etype, exc, tb):
                 # TODO: check if PEP558 changes the need for this call
                 # https://github.com/python/cpython/pull/3640
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         if self._cancel_status is not None:
             binding = "active"
         elif self._has_been_entered:
@@ -634,7 +642,7 @@ def __repr__(self):
 
     @contextmanager
     @enable_ki_protection
-    def _might_change_registered_deadline(self):
+    def _might_change_registered_deadline(self) -> Iterator[None]:
         try:
             yield
         finally:
@@ -658,7 +666,7 @@ def _might_change_registered_deadline(self):
                         runner.force_guest_tick_asap()
 
     @property
-    def deadline(self):
+    def deadline(self) -> float:
         """Read-write, :class:`float`. An absolute time on the current
         run's clock at which this scope will automatically become
         cancelled. You can adjust the deadline by modifying this
@@ -684,12 +692,12 @@ def deadline(self):
         return self._deadline
 
     @deadline.setter
-    def deadline(self, new_deadline):
+    def deadline(self, new_deadline: float) -> None:
         with self._might_change_registered_deadline():
             self._deadline = float(new_deadline)
 
     @property
-    def shield(self):
+    def shield(self) -> bool:
         """Read-write, :class:`bool`, default :data:`False`. So long as
         this is set to :data:`True`, then the code inside this scope
         will not receive :exc:`~trio.Cancelled` exceptions from scopes
@@ -714,15 +722,15 @@ def shield(self):
 
     @shield.setter
     @enable_ki_protection
-    def shield(self, new_value):
+    def shield(self, new_value: bool) -> None:
         if not isinstance(new_value, bool):
             raise TypeError("shield must be a bool")
         self._shield = new_value
         if self._cancel_status is not None:
             self._cancel_status.recalculate()
 
     @enable_ki_protection
-    def cancel(self):
+    def cancel(self) -> None:
         """Cancels this scope immediately.
 
         This method is idempotent, i.e., if the scope was already
@@ -736,7 +744,7 @@ def cancel(self):
             self._cancel_status.recalculate()
 
     @property
-    def cancel_called(self):
+    def cancel_called(self) -> bool:
         """Readonly :class:`bool`. Records whether cancellation has been
         requested for this scope, either by an explicit call to
         :meth:`cancel` or by the deadline expiring.
@@ -890,7 +898,9 @@ def __exit__(self):  # pragma: no cover
         assert False, """Never called, but should be defined"""
 
 
-def open_nursery(strict_exception_groups=None):
+def open_nursery(
+    strict_exception_groups: bool | None = None,
+) -> AbstractAsyncContextManager[Nursery]:
     """Returns an async context manager which must be used to create a
     new `Nursery`.
 

trio/_tests/check_type_completeness.py

@@ -75,6 +75,19 @@ def main(args: argparse.Namespace) -> int:
     if res.stderr:
         print(res.stderr)
 
+    if args.full_diagnostics_file is not None:
+        with open(args.full_diagnostics_file, "w") as file:
+            json.dump(
+                [
+                    sym
+                    for sym in current_result["typeCompleteness"]["symbols"]
+                    if sym["diagnostics"]
+                ],
+                file,
+                sort_keys=True,
+                indent=2,
+            )
+
     last_result = json.loads(RESULT_FILE.read_text())
 
     for key in "errorCount", "warningCount", "informationCount":
@@ -153,6 +166,7 @@ def main(args: argparse.Namespace) -> int:
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--overwrite-file", action="store_true", default=False)
+parser.add_argument("--full-diagnostics-file", type=Path, default=None)
 args = parser.parse_args()
 
 assert __name__ == "__main__", "This script should be run standalone"

trio/_tests/verify_types.json

@@ -7,11 +7,11 @@
     "warningCount": 0
   },
   "typeCompleteness": {
-    "completenessScore": 0.8155339805825242,
+    "completenessScore": 0.8317152103559871,
     "exportedSymbolCounts": {
       "withAmbiguousType": 1,
-      "withKnownType": 504,
-      "withUnknownType": 113
+      "withKnownType": 514,
+      "withUnknownType": 103
     },
     "ignoreUnknownTypesFromImports": true,
     "missingClassDocStringCount": 1,
@@ -45,54 +45,32 @@
       }
     ],
     "otherSymbolCounts": {
-      "withAmbiguousType": 15,
-      "withKnownType": 231,
-      "withUnknownType": 236
+      "withAmbiguousType": 14,
+      "withKnownType": 244,
+      "withUnknownType": 224
     },
     "packageName": "trio",
     "symbols": [
-      "trio._core._exceptions.Cancelled",
-      "trio._core._exceptions.Cancelled.__str__",
-      "trio._util.NoPublicConstructor",
-      "trio._util.NoPublicConstructor.__call__",
-      "trio._util.Final.__new__",
       "trio.run",
-      "trio.open_nursery",
-      "trio._core._run.CancelScope",
-      "trio._core._run.CancelScope.cancelled_caught",
-      "trio._core._run.CancelScope.__exit__",
-      "trio._core._run.CancelScope.__repr__",
-      "trio._core._run.CancelScope.deadline",
-      "trio._core._run.CancelScope.cancel_called",
       "trio.current_effective_deadline",
       "trio._core._run._TaskStatusIgnored.__repr__",
       "trio._core._run._TaskStatusIgnored.started",
       "trio.current_time",
-      "trio._core._run.Nursery",
       "trio._core._run.Nursery.__init__",
       "trio._core._run.Nursery.child_tasks",
       "trio._core._run.Nursery.parent_task",
       "trio._core._run.Nursery.start_soon",
       "trio._core._run.Nursery.start",
       "trio._core._run.Nursery.__del__",
-      "trio.move_on_at",
-      "trio.move_on_after",
-      "trio.sleep_forever",
-      "trio.sleep_until",
-      "trio.sleep",
-      "trio.fail_after",
-      "trio._sync.Event",
       "trio._sync.Event.is_set",
       "trio._sync.Event.wait",
       "trio._sync.Event.statistics",
-      "trio._sync.CapacityLimiter",
       "trio._sync.CapacityLimiter.__init__",
       "trio._sync.CapacityLimiter.__repr__",
       "trio._sync.CapacityLimiter.total_tokens",
       "trio._sync.CapacityLimiter.borrowed_tokens",
       "trio._sync.CapacityLimiter.available_tokens",
       "trio._sync.CapacityLimiter.statistics",
-      "trio._sync.Semaphore",
       "trio._sync.Semaphore.__init__",
       "trio._sync.Semaphore.__repr__",
       "trio._sync.Semaphore.value",
@@ -103,7 +81,6 @@
       "trio._sync._LockImpl.locked",
       "trio._sync._LockImpl.statistics",
       "trio._sync.StrictFIFOLock",
-      "trio._sync.Condition",
       "trio._sync.Condition.__init__",
       "trio._sync.Condition.locked",
       "trio._sync.Condition.acquire_nowait",
@@ -165,7 +142,6 @@
       "trio._path.Path.__bytes__",
       "trio._path.Path.__truediv__",
       "trio._path.Path.__rtruediv__",
-      "trio._path.AsyncAutoWrapperType",
       "trio._path.AsyncAutoWrapperType.__init__",
       "trio._path.AsyncAutoWrapperType.generate_forwards",
       "trio._path.AsyncAutoWrapperType.generate_wraps",
@@ -203,7 +179,6 @@
       "trio._ssl.SSLListener.__init__",
       "trio._ssl.SSLListener.accept",
       "trio._ssl.SSLListener.aclose",
-      "trio._dtls.DTLSEndpoint",
       "trio._dtls.DTLSEndpoint.__init__",
       "trio._dtls.DTLSEndpoint.__del__",
       "trio._dtls.DTLSEndpoint.close",
@@ -254,7 +229,6 @@
       "trio.from_thread.run_sync",
       "trio.lowlevel.cancel_shielded_checkpoint",
       "trio.lowlevel.currently_ki_protected",
-      "trio._core._run.Task",
       "trio._core._run.Task.coro",
       "trio._core._run.Task.name",
       "trio._core._run.Task.context",
@@ -266,13 +240,11 @@
       "trio._core._run.Task.iter_await_frames",
       "trio.lowlevel.checkpoint",
       "trio.lowlevel.current_task",
-      "trio._core._parking_lot.ParkingLot",
       "trio._core._parking_lot.ParkingLot.__len__",
       "trio._core._parking_lot.ParkingLot.__bool__",
       "trio._core._parking_lot.ParkingLot.unpark_all",
       "trio._core._parking_lot.ParkingLot.repark_all",
       "trio._core._parking_lot.ParkingLot.statistics",
-      "trio._core._unbounded_queue.UnboundedQueue",
       "trio._core._unbounded_queue.UnboundedQueue.__repr__",
       "trio._core._unbounded_queue.UnboundedQueue.qsize",
       "trio._core._unbounded_queue.UnboundedQueue.empty",
@@ -281,12 +253,10 @@
       "trio._core._unbounded_queue.UnboundedQueue.statistics",
       "trio._core._unbounded_queue.UnboundedQueue.__aiter__",
       "trio._core._unbounded_queue.UnboundedQueue.__anext__",
-      "trio._core._local.RunVar",
       "trio._core._local.RunVar.get",
       "trio._core._local.RunVar.set",
       "trio._core._local.RunVar.reset",
       "trio._core._local.RunVar.__repr__",
-      "trio._core._entry_queue.TrioToken",
       "trio._core._entry_queue.TrioToken.run_sync_soon",
       "trio.lowlevel.current_trio_token",
       "trio.lowlevel.temporarily_detach_coroutine_object",
@@ -329,7 +299,6 @@
       "trio.testing.trio_test",
       "trio.testing.assert_checkpoints",
       "trio.testing.assert_no_checkpoints",
-      "trio.testing._sequencer.Sequencer",
       "trio.testing.check_one_way_stream",
       "trio.testing.check_two_way_stream",
       "trio.testing.check_half_closeable_stream",

trio/_timeouts.py

@@ -1,10 +1,13 @@
+from __future__ import annotations
+
 import math
-from contextlib import contextmanager
+from contextlib import AbstractContextManager, contextmanager
+from typing import TYPE_CHECKING
 
 import trio
 
 
-def move_on_at(deadline):
+def move_on_at(deadline: float) -> trio.CancelScope:
     """Use as a context manager to create a cancel scope with the given
     absolute deadline.
 
@@ -20,7 +23,7 @@ def move_on_at(deadline):
     return trio.CancelScope(deadline=deadline)
 
 
-def move_on_after(seconds):
+def move_on_after(seconds: float) -> trio.CancelScope:
     """Use as a context manager to create a cancel scope whose deadline is
     set to now + *seconds*.
 
@@ -36,7 +39,7 @@ def move_on_after(seconds):
     return move_on_at(trio.current_time() + seconds)
 
 
-async def sleep_forever():
+async def sleep_forever() -> None:
     """Pause execution of the current task forever (or until cancelled).
 
     Equivalent to calling ``await sleep(math.inf)``.
@@ -45,7 +48,7 @@ async def sleep_forever():
     await trio.lowlevel.wait_task_rescheduled(lambda _: trio.lowlevel.Abort.SUCCEEDED)
 
 
-async def sleep_until(deadline):
+async def sleep_until(deadline: float) -> None:
     """Pause execution of the current task until the given time.
 
     The difference between :func:`sleep` and :func:`sleep_until` is that the
@@ -65,7 +68,7 @@ async def sleep_until(deadline):
         await sleep_forever()
 
 
-async def sleep(seconds):
+async def sleep(seconds: float) -> None:
     """Pause execution of the current task for the given number of seconds.
 
     Args:
@@ -91,8 +94,9 @@ class TooSlowError(Exception):
     """
 
 
-@contextmanager
-def fail_at(deadline):
+# workaround for PyCharm not being able to infer return type from @contextmanager
+# see https://youtrack.jetbrains.com/issue/PY-36444/PyCharm-doesnt-infer-types-when-using-contextlib.contextmanager-decorator
+def fail_at(deadline: float) -> AbstractContextManager[trio.CancelScope]:  # type: ignore[misc]
     """Creates a cancel scope with the given deadline, and raises an error if it
     is actually cancelled.
 
@@ -113,14 +117,17 @@ def fail_at(deadline):
       ValueError: if deadline is NaN.
 
     """
-
     with move_on_at(deadline) as scope:
         yield scope
     if scope.cancelled_caught:
         raise TooSlowError
 
 
-def fail_after(seconds):
+if not TYPE_CHECKING:
+    fail_at = contextmanager(fail_at)
+
+
+def fail_after(seconds: float) -> AbstractContextManager[trio.CancelScope]:
     """Creates a cancel scope with the given timeout, and raises an error if
     it is actually cancelled.