sys.set_asyncgen_hook by youknowone · Pull Request #6439 · RustPython/RustPython

youknowone

Walkthrough

The changes introduce lifecycle and state management patterns to async generators and coroutines, adding hook initialization, finalizer invocation, and closed-state tracking to prevent reuse after completion. Additionally, await() result validation is added to enforce iterator protocol compliance.

Changes

Cohort / File(s)	Summary
Async Generator Lifecycle Management `crates/vm/src/builtins/asyncgenerator.rs`	Adds hook lifecycle tracking (`ag_hooks_inited`, `ag_finalizer` fields) with `init_hooks()` method to read thread-local finalizer and firstiter hooks. Introduces `call_finalizer()` for guarded finalizer invocation. Updates method signatures (`__anext__`, `asend`, `athrow`, `aclose`) to return `PyResult` and initialize hooks before returning awaitables.
Awaitable State Management `crates/vm/src/builtins/asyncgenerator.rs`	Adds `state` field to `PyAnextAwaitable` with `AwaitableState` enum for lifecycle tracking. Implements `check_closed()` validation; `send()` and `throw()` now call it and update state to `Iter`, while `close()` marks state as `Closed`.
Coroutine State Protection `crates/vm/src/builtins/coroutine.rs`	Adds `closed: AtomicCell<bool>` field to `PyCoroutineWrapper` for reuse prevention. Implements `check_closed()` helper. Methods `send()`, `throw()`, and `close()` now track closed state, with `send()` and `throw()` marking closed on `StopIteration`.
Awaitable Result Validation `crates/vm/src/frame.rs`	Adds validation in `GetAwaitable` handling to ensure `__await__()` returns an iterator; raises `TypeError` with message "await() returned non-iterator of type 'X'" if violated.

Sequence Diagram

sequenceDiagram
    participant App as Application
    participant VM as VirtualMachine
    participant AG as AsyncGenerator
    participant Hooks as Thread-local Hooks
    participant Awaitable as PyAnextAwaitable
    
    App->>AG: Call __anext__()
    activate AG
    AG->>VM: init_hooks(self, vm)
    activate VM
    VM->>Hooks: Read finalizer & firstiter
    Hooks-->>VM: Return hooks
    VM->>AG: Store finalizer in ag_finalizer
    VM->>VM: Call firstiter (if present)
    VM-->>AG: Hooks initialized
    deactivate VM
    AG->>Awaitable: Create new PyAnextAwaitable(state=Init)
    AG-->>App: Return Awaitable
    deactivate AG
    
    App->>Awaitable: Call send(val)
    activate Awaitable
    Awaitable->>Awaitable: check_closed()?
    alt Not Closed
        Awaitable->>Awaitable: state = Iter
        Awaitable-->>App: Yield value
    else Closed
        Awaitable-->>App: RuntimeError
    end
    deactivate Awaitable
    
    App->>Awaitable: Call close()
    activate Awaitable
    Awaitable->>Awaitable: state = Closed
    Awaitable-->>App: Closed
    deactivate Awaitable
    
    Note over AG,Awaitable: On AsyncGen finalization
    AG->>AG: call_finalizer(self, vm)
    AG->>VM: Invoke stored finalizer

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

High-complexity areas: asyncgenerator.rs — dense state management with hook lifecycle, finalizer handling, and multiple method signature changes; requires understanding of thread-local hook semantics and state transitions across PyAnextAwaitable.
Moderate-complexity areas: coroutine.rs — similar state pattern but narrower scope; validate consistency with async generator approach.
Attention needed: Verify init_hooks() thread-safety and finalizer invocation semantics; ensure closed-state checks prevent all reuse paths; validate __await__() result validation in frame.rs against protocol requirements.

Possibly related PRs

PyAnextAwaitable #6427 — Introduces the PyAnextAwaitable type and qualname support for async generators; this PR extends that work by adding lifecycle hooks, finalizer handling, and state management to the same types.

Poem

🐰 Hooks and finalizers now take flight,
Async generators locked up tight,
Closed-state checks prevent reuse with care,
Lifecycle management everywhere!
The awaitable dance—once, then done.

Pre-merge checks and finishing touches

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The PR title 'sys.set_asyncgen_hook' directly references the main feature being implemented, which is a system hook for async generators based on the code changes introducing async generator hook lifecycle management.
Docstring Coverage	✅ Passed	Docstring coverage is 90.91% which is sufficient. The required threshold is 80.00%.

✨ Finishing touches

📝 Generate docstrings

🧪 Generate unit tests (beta)

Create PR with unit tests
Post copyable unit tests in a comment

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands and usage tips.}