Add typing for _core/_asyncgens.py

pyproject.toml

@@ -50,14 +50,15 @@ disallow_untyped_defs = false
 # Fully typed, enable stricter checks
 module = [
     "trio._abc",
+    "trio._core._asyncgens",
     "trio._core._entry_queue",
     "trio._core._local",
     "trio._core._unbounded_queue",
     "trio._core._thread_cache",
     "trio._deprecate",
     "trio._dtls",
     "trio._file_io",
-    "trio._highlevel_open_tcp_stream.py",
+    "trio._highlevel_open_tcp_stream",
     "trio._ki",
     "trio._socket",
     "trio._sync",

trio/_core/_asyncgens.py

@@ -1,7 +1,11 @@
+from __future__ import annotations
+
 import logging
 import sys
 import warnings
 import weakref
+from types import AsyncGeneratorType
+from typing import TYPE_CHECKING, NoReturn
 
 import attr
 
@@ -12,6 +16,15 @@
 # Used to log exceptions in async generator finalizers
 ASYNCGEN_LOGGER = logging.getLogger("trio.async_generator_errors")
 
+if TYPE_CHECKING:
+    from typing import Set
+
+    _WEAK_ASYNC_GEN_SET = weakref.WeakSet[AsyncGeneratorType[object, NoReturn]]
+    _ASYNC_GEN_SET = Set[AsyncGeneratorType[object, NoReturn]]
+else:
+    _WEAK_ASYNC_GEN_SET = weakref.WeakSet
+    _ASYNC_GEN_SET = set
+
 
 @attr.s(eq=False, slots=True)
 class AsyncGenerators:
@@ -22,17 +35,17 @@ class AsyncGenerators:
     # asyncgens after the system nursery has been closed, it's a
     # regular set so we don't have to deal with GC firing at
     # unexpected times.
-    alive = attr.ib(factory=weakref.WeakSet)
+    alive: _WEAK_ASYNC_GEN_SET | _ASYNC_GEN_SET = attr.ib(factory=_WEAK_ASYNC_GEN_SET)
 
     # This collects async generators that get garbage collected during
     # the one-tick window between the system nursery closing and the
     # init task starting end-of-run asyncgen finalization.
-    trailing_needs_finalize = attr.ib(factory=set)
+    trailing_needs_finalize: _ASYNC_GEN_SET = attr.ib(factory=_ASYNC_GEN_SET)
 
     prev_hooks = attr.ib(init=False)
 
-    def install_hooks(self, runner):
-        def firstiter(agen):
+    def install_hooks(self, runner: _run.Runner) -> None:
+        def firstiter(agen: AsyncGeneratorType[object, NoReturn]) -> None:
             if hasattr(_run.GLOBAL_RUN_CONTEXT, "task"):
                 self.alive.add(agen)
             else:
@@ -46,7 +59,9 @@ def firstiter(agen):
                 if self.prev_hooks.firstiter is not None:
                     self.prev_hooks.firstiter(agen)
 
-        def finalize_in_trio_context(agen, agen_name):
+        def finalize_in_trio_context(
+            agen: AsyncGeneratorType[object, NoReturn], agen_name: str
+        ) -> None:
             try:
                 runner.spawn_system_task(
                     self._finalize_one,
@@ -61,7 +76,7 @@ def finalize_in_trio_context(agen, agen_name):
                 # have hit it.
                 self.trailing_needs_finalize.add(agen)
 
-        def finalizer(agen):
+        def finalizer(agen: AsyncGeneratorType[object, NoReturn]) -> None:
             agen_name = name_asyncgen(agen)
             try:
                 is_ours = not agen.ag_frame.f_locals.get("@trio_foreign_asyncgen")
@@ -112,9 +127,9 @@ def finalizer(agen):
                         )
 
         self.prev_hooks = sys.get_asyncgen_hooks()
-        sys.set_asyncgen_hooks(firstiter=firstiter, finalizer=finalizer)
+        sys.set_asyncgen_hooks(firstiter=firstiter, finalizer=finalizer)  # type: ignore[arg-type]  # Finalizer doesn't use AsyncGeneratorType
 
-    async def finalize_remaining(self, runner):
+    async def finalize_remaining(self, runner: _run.Runner) -> None:
         # This is called from init after shutting down the system nursery.
         # The only tasks running at this point are init and
         # the run_sync_soon task, and since the system nursery is closed,
@@ -170,14 +185,16 @@ async def finalize_remaining(self, runner):
         # all are gone.
         while self.alive:
             batch = self.alive
-            self.alive = set()
+            self.alive = _ASYNC_GEN_SET()
             for agen in batch:
                 await self._finalize_one(agen, name_asyncgen(agen))
 
-    def close(self):
+    def close(self) -> None:
         sys.set_asyncgen_hooks(*self.prev_hooks)
 
-    async def _finalize_one(self, agen, name):
+    async def _finalize_one(
+        self, agen: AsyncGeneratorType[object, NoReturn], name: object
+    ) -> None:
         try:
             # This shield ensures that finalize_asyncgen never exits
             # with an exception, not even a Cancelled. The inside