From dd10a5290cfbe2c6d529c5b7f706bf52965445cb Mon Sep 17 00:00:00 2001 From: David Lamparter Date: Fri, 6 Sep 2024 17:08:05 +0200 Subject: [PATCH] topotato: don't use GC to track generators Much faster to just keep a list. Signed-off-by: David Lamparter --- selftests/test_generatorwrap.py | 17 ++++- topotato/generatorwrap.py | 118 ++++++++++---------------------- 2 files changed, 52 insertions(+), 83 deletions(-) diff --git a/selftests/test_generatorwrap.py b/selftests/test_generatorwrap.py index 6826419..42098de 100644 --- a/selftests/test_generatorwrap.py +++ b/selftests/test_generatorwrap.py @@ -6,7 +6,12 @@ """ import pytest -from topotato.generatorwrap import GeneratorChecks, GeneratorWrapper, GeneratorsUnused +from topotato.generatorwrap import ( + GeneratorChecks, + GeneratorWrapper, + GeneratorsUnused, + GeneratorChecksWarning, +) def ref_gen_nowrap(): @@ -23,10 +28,17 @@ def both_gens(request): return request.param +def test_no_context_warn(): + with pytest.warns(GeneratorChecksWarning): + assert list(ref_gen_wrap()) == [1, 2, 3] + + +@pytest.mark.filterwarnings("ignore::topotato.generatorwrap.GeneratorChecksWarning") def test_still_functional(both_gens): assert list(both_gens()) == [1, 2, 3] +@pytest.mark.filterwarnings("ignore::topotato.generatorwrap.GeneratorChecksWarning") def test_sending(both_gens): gen = both_gens() itr = iter(gen) @@ -59,7 +71,7 @@ def test_wrapped_multifail(): with GeneratorChecks() as checks: ref_gen_wrap() ref_gen_wrap() - assert len(checks._errs) == 2 + assert len(checks._active) == 2 class GeneratorMethod: @@ -78,6 +90,7 @@ def gen3(): yield 3 +@pytest.mark.filterwarnings("ignore::topotato.generatorwrap.GeneratorChecksWarning") def test_method_wrap(): assert list(GeneratorMethod().gen1()) == [1] assert list(GeneratorMethod.gen2()) == [2] diff --git a/topotato/generatorwrap.py b/topotato/generatorwrap.py index c1e0e3d..cf5685d 100644 --- a/topotato/generatorwrap.py +++ b/topotato/generatorwrap.py @@ -5,14 +5,12 @@ Wrap generator functions & raise exceptions if a generator is never used. """ -import sys import functools import traceback import inspect -import gc +import warnings from types import TracebackType from typing import ( - Any, Callable, Generator, Generic, @@ -23,26 +21,8 @@ ) -class GeneratorDeletedUnused(Exception): - """ - Exception raised from :py:meth:`GeneratorWrapper.__del__` if the generator - never actually ran. - - Note that exceptions raised from __del__ are handled through - sys.unraisablehook since __del__ is called asynchronously and exceptions - are normally printed but ignored there. - """ - - origin: List[str] - """ - Formatted string traceback for the generator's call location. - - Not saving FrameInfo/related here since the frame continues executing - after the call site and these objects would change due to that. - """ - - def __init__(self, origin: List[str]): - self.origin = origin +class GeneratorChecksWarning(UserWarning): + pass class GeneratorsUnused(Exception): @@ -51,56 +31,38 @@ class GeneratorsUnused(Exception): :py:class:`GeneratorChecks` ``with`` context. """ - exceptions: List[GeneratorDeletedUnused] + generators: List[Generator] - def __init__(self, exceptions: List[GeneratorDeletedUnused]): + def __init__(self, generators: List[Generator]): super().__init__() - self.exceptions = exceptions + self.generators = generators def __str__(self): - items = "\n===\n".join( - e.origin[-1].rstrip("\n ") for e in self.exceptions - ).replace("\n", "\n ") - return f'{len(self.exceptions)} generators were invoked but never executed (forgotten "yield from"?):\n {items}' + items = "\n===\n".join(g._loc[-1] for g in self.generators).replace( + "\n", "\n " + ) + return f'{len(self.generators)} generators were invoked but never executed (forgotten "yield from"?):\n {items}' class GeneratorChecks: """ - Context manager to collect raised :py:class:`GeneratorDeletedUnused`. - - Since __del__ goes to sys.unraisablehook, this context manager places - itself into that hook for the duration of its ``with`` block. All - GeneratorDeletedUnused exceptions are collected and then raised at the - end of the context. + Context manager to track called but not iterated GeneratorWrapper. """ - _prev_hook: Optional[Callable[[Any], Any]] = None - """ - Original sys.unraisablehook to restore after. - """ - _errs: List[GeneratorDeletedUnused] + _active: List """ - Accumulated exceptions, if any. + Calling a GeneratorWrapper-wrapped iterator adds it here, iterating over it + removes it. At the end of the context, this list needs to be empty or we + raise an exception. """ - def __init__(self): - self._errs = [] + _ctxstack: List["GeneratorChecks"] = [] - def _unraisable(self, args) -> None: - """ - Handler to be installed into sys.unraisablehook. - """ - assert self._prev_hook - - if isinstance(args.exc_value, GeneratorDeletedUnused): - self._errs.append(args.exc_value) - return - - self._prev_hook(args) + def __init__(self): + self._active = [] def __enter__(self): - self._prev_hook = sys.unraisablehook - sys.unraisablehook = self._unraisable + GeneratorChecks._ctxstack.insert(0, self) return self def __exit__( @@ -109,19 +71,13 @@ def __exit__( exc_value: Optional[BaseException], tb: Optional[TracebackType], ): - assert self._prev_hook - - gc.collect() - sys.unraisablehook = self._prev_hook - self._prev_hook = None - if exc_type is not None: return - if not self._errs: - return + assert GeneratorChecks._ctxstack.pop(0) == self - raise GeneratorsUnused(self._errs) + if self._active: + raise GeneratorsUnused(self._active) TG = TypeVar("TG", bound=Generator) @@ -138,7 +94,7 @@ def my_generator(foo): yield foo - 1 When the generator is later _called_ but not actually iterated over, it - will raise :py:class:`GeneratorDeletedUnused` from its __del__ method. + will register itself on the innermost :py:class:`GeneratorChecks` context. These should be collected like this:: with GeneratorChecks(): @@ -154,32 +110,32 @@ def my_generator(foo): """ Original generator object to forward __iter__ to. """ - _run: bool - """ - Has this generator actually been iterated over? (Not necessarily to completion.) - """ _loc: List[str] """ Location of call to pass to :py:class:`GeneratorDeletedUnused`, see there. """ + _gchk: Optional[GeneratorChecks] def __init__(self, wraps: TG, loc: List[str]): self._wraps = wraps self._loc = loc - self._run = False + + if GeneratorChecks._ctxstack: + self._gchk = GeneratorChecks._ctxstack[0] + self._gchk._active.append(self) + else: + self._gchk = None + warnings.warn( + "GeneratorWrapper-wrapped generator invoked without active GeneratorChecks", + GeneratorChecksWarning, + ) def __iter__(self): - self._run = True - del self._loc + if self._gchk: + self._gchk._active.remove(self) + self._gchk = None return self._wraps.__iter__() - def __del__(self): - if self._run: - return - if isinstance(self._wraps, GeneratorWrapper): - self._wraps._run = True - raise GeneratorDeletedUnused(self._loc) - @classmethod def apply(cls, function: Callable[..., TG]) -> Callable[..., TG]: """