diff --git a/src/trio/_core/_asyncgens.py b/src/trio/_core/_asyncgens.py index 1ec7b6471..89976d25e 100644 --- a/src/trio/_core/_asyncgens.py +++ b/src/trio/_core/_asyncgens.py @@ -25,6 +25,32 @@ _ASYNC_GEN_SET = set +@_core.disable_ki_protection +def _finalize_without_ki_protection( + agen_name: str, + agen: AsyncGeneratorType[object, NoReturn], +) -> None: + # Host has no finalizer. Reimplement the default + # Python behavior with no hooks installed: throw in + # GeneratorExit, step once, raise RuntimeError if + # it doesn't exit. + closer = agen.aclose() + try: + # If the next thing is a yield, this will raise RuntimeError + # which we allow to propagate + closer.send(None) + except StopIteration: + pass + else: + # If the next thing is an await, we get here. Give a nicer + # error than the default "async generator ignored GeneratorExit" + raise RuntimeError( + f"Non-Trio async generator {agen_name!r} awaited something " + "during finalization; install a finalization hook to " + "support this, or wrap it in 'async with aclosing(...):'", + ) + + @attrs.define(eq=False) class AsyncGenerators: # Async generators are added to this set when first iterated. Any @@ -78,6 +104,7 @@ def finalize_in_trio_context( # have hit it. self.trailing_needs_finalize.add(agen) + @_core.enable_ki_protection def finalizer(agen: AsyncGeneratorType[object, NoReturn]) -> None: try: self.foreign.remove(id(agen)) @@ -112,25 +139,7 @@ def finalizer(agen: AsyncGeneratorType[object, NoReturn]) -> None: if self.prev_hooks.finalizer is not None: self.prev_hooks.finalizer(agen) else: - # Host has no finalizer. Reimplement the default - # Python behavior with no hooks installed: throw in - # GeneratorExit, step once, raise RuntimeError if - # it doesn't exit. - closer = agen.aclose() - try: - # If the next thing is a yield, this will raise RuntimeError - # which we allow to propagate - closer.send(None) - except StopIteration: - pass - else: - # If the next thing is an await, we get here. Give a nicer - # error than the default "async generator ignored GeneratorExit" - raise RuntimeError( - f"Non-Trio async generator {agen_name!r} awaited something " - "during finalization; install a finalization hook to " - "support this, or wrap it in 'async with aclosing(...):'", - ) + _finalize_without_ki_protection(agen_name, agen) self.prev_hooks = sys.get_asyncgen_hooks() sys.set_asyncgen_hooks(firstiter=firstiter, finalizer=finalizer) # type: ignore[arg-type] # Finalizer doesn't use AsyncGeneratorType