diff --git a/pyproject.toml b/pyproject.toml index a23f7f5db9..614b2f7c5e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,6 +50,7 @@ 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", @@ -57,7 +58,7 @@ module = [ "trio._deprecate", "trio._dtls", "trio._file_io", - "trio._highlevel_open_tcp_stream.py", + "trio._highlevel_open_tcp_stream", "trio._ki", "trio._socket", "trio._sync", diff --git a/trio/_core/_asyncgens.py b/trio/_core/_asyncgens.py index 5f02ebe76d..4261328278 100644 --- a/trio/_core/_asyncgens.py +++ b/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