fix False positive with overloaded async iterators #2873

[asukaminato0721]

Summary

Fixes #2873

The change has two parts.

overload signatures are normalized when the implementation is an async generator, so async def overload stubs returning AsyncIterator[...] no longer stay wrapped as Coroutine[..., AsyncIterator[...]].

generator decomposition now handles unions member-by-member, which fixes the bogus yield Patch() vs FullState error from AsyncIterator[Patch] | AsyncIterator[FullState].

Test Plan

add test

[Copilot]

Pull request overview

Fixes a false-positive overload/yield error involving async generator implementations with overload stubs returning AsyncIterator[...], and unions of async-iterator return types.

Changes:

• Normalize overload return signatures when the implementation is an async generator (so overloads don’t remain wrapped as Coroutine[..., AsyncIterator[...]]).
• Update generator/async-generator decomposition to handle Union types member-by-member and union the decomposed components.
• Add a regression test covering overloaded async iterator returns and yield behavior.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
pyrefly/lib/test/overload.rs	Adds a regression test reproducing issue #2873 with overloaded async generator returning AsyncIterator[...].
pyrefly/lib/alt/unwrap.rs	Decomposes generator / async generator types across union members and unions their yield/send/return components.
pyrefly/lib/alt/function.rs	Normalizes overload signatures for async generator implementations before overload consistency checking.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[github-actions]

Diff from mypy_primer, showing the effect of this PR on open source code:

freqtrade (https://github.com/freqtrade/freqtrade)
+ ERROR freqtrade/rpc/api_server/deps.py:23:24-49: Async generator function should return `AsyncGenerator` [bad-return]

starlette (https://github.com/encode/starlette)
+ ERROR tests/middleware/test_gzip.py:75:58-71: Async generator function should return `AsyncGenerator` [bad-return]
+ ERROR tests/middleware/test_gzip.py:98:58-71: Async generator function should return `AsyncGenerator` [bad-return]
+ ERROR tests/middleware/test_gzip.py:123:58-71: Async generator function should return `AsyncGenerator` [bad-return]
+ ERROR tests/middleware/test_gzip.py:146:58-71: Async generator function should return `AsyncGenerator` [bad-return]

steam.py (https://github.com/Gobot1234/steam.py)
- ERROR steam/leaderboard.py:101:15-22: Overload return type `Coroutine[Unknown, Unknown, AsyncGenerator[LeaderboardUser]]` is not assignable to implementation return type `AsyncGenerator[LeaderboardUser]` [inconsistent-overload]
- ERROR steam/leaderboard.py:110:15-22: Overload return type `Coroutine[Unknown, Unknown, AsyncGenerator[LeaderboardUser]]` is not assignable to implementation return type `AsyncGenerator[LeaderboardUser]` [inconsistent-overload]

strawberry (https://github.com/strawberry-graphql/strawberry)
+ ERROR strawberry/channels/testing.py:116:10-58: Async generator function should return `AsyncGenerator` [bad-return]

optuna (https://github.com/optuna/optuna)
- ERROR optuna/storages/_rdb/alembic/versions/v2.4.0.a.py:181:42-51: Argument `str | None` is not assignable to parameter `constraint_name` with type `str` in function `alembic.operations.base.BatchOperations.drop_constraint` [bad-argument-type]

prefect (https://github.com/PrefectHQ/prefect)
+ ERROR src/prefect/server/models/task_runs.py:138:8-21: Object of class `NoneType` has no attribute `created` [missing-attribute]
+ ERROR src/prefect/server/models/task_runs.py:141:25-33: Object of class `NoneType` has no attribute `id` [missing-attribute]
+ ERROR src/prefect/server/models/task_runs.py:147:12-17: Returned type `TaskRun | None` is not assignable to declared return type `TaskRun` [bad-return]

[github-actions]

Primer Diff Classification

❌ 3 regression(s) | ✅ 3 improvement(s) | 6 project(s) total | +9, -3 errors

3 regression(s) across freqtrade, optuna, prefect. error kinds: bad-return, bad-argument-type, missing-attribute. caused by decompose_async_generator(), decompose_generator(). 3 improvement(s) across starlette, steam.py, strawberry.

Project	Verdict	Changes	Error Kinds	Root Cause
freqtrade	❌ Regression	+1	bad-return	decompose_async_generator()
starlette	✅ Improvement	+4	bad-return	decompose_async_generator()
steam.py	✅ Improvement	-2	inconsistent-overload	normalize_async_generator_overloads()
strawberry	✅ Improvement	+1	bad-return	decompose_async_generator()
optuna	❌ Regression	-1	bad-argument-type	decompose_generator()
prefect	❌ Regression	+3	bad-return, missing-attribute	decompose_generator()

Detailed analysis

❌ Regression (3)

freqtrade (+1)

The function get_rpc() is an async generator (it contains yield _rpc). Its return type is annotated as AsyncIterator[RPC] | None. The bad-return error states that an async generator function should return AsyncGenerator.

The core issue is that the return type annotation AsyncIterator[RPC] | None is not a standard return annotation for an async generator function. An async generator should be annotated with AsyncIterator[YieldType] or AsyncGenerator[YieldType, SendType], not a union of these with None.

Looking at the code: the else branch raises an exception rather than returning None, so the | None in the return type is unnecessary and misleading. The function always either yields values (async generator path) or raises an exception.

However, this is likely a true positive that is low-severity in practice. The annotation AsyncIterator[RPC] | None is technically incorrect for an async generator - the | None part is meaningless since async generators don't return values in the traditional sense. Other type checkers (mypy, pyright) happen to be lenient about this and don't flag it, but pyrefly's stricter checking of the return type annotation for async generators is not unreasonable.

That said, since mypy and pyright both accept this pattern without error, and it appears in real-world codebases like freqtrade, this represents a compatibility difference where pyrefly is stricter than other type checkers. Whether this is a "false positive" depends on perspective - the annotation is technically wrong, but the intent is clear and other tools accept it.

This is best categorized as a pyrefly-specific strictness issue where pyrefly does not handle union return types containing a valid async generator type (like AsyncIterator[X]) combined with None.

Attribution: The change to decompose_async_generator() in pyrefly/lib/alt/unwrap.rs introduced member-by-member union decomposition. When encountering AsyncIterator[RPC] | None, it tries to decompose None as an async generator, fails, and the collect() on Option<Vec<_>> causes the entire decomposition to return None. This triggers the bad-return error. The old code handled the union as a whole via is_subset_eq, which worked correctly.

optuna (-1)

The removed error was a true positive — c['name'] genuinely has type str | None (SQLAlchemy's get_unique_constraints can return None for constraint names), and drop_constraint requires str. Removing detection of this real type mismatch is a regression, even though it appears to be an unintended side effect of the async generator union handling changes.
Attribution: The changes to decompose_generator() and decompose_async_generator() in pyrefly/lib/alt/unwrap.rs refactored how unions are handled member-by-member. This may have indirectly affected type inference for dict value lookups, though the connection to this specific str | None error is not obvious.

prefect (+3)

These are false positives caused by pyrefly failing to narrow model: Union[orm_models.TaskRun, None] to orm_models.TaskRun after an if/else block where both branches guarantee a non-None value by line 136. In the if task_run.flow_run_id branch (line 108), scalar_one() returns a non-optional value (it raises NoResultFound if no row exists, so its return type is orm_models.TaskRun). In the else branch, result.scalar() on line 125 returns Optional[orm_models.TaskRun], but the subsequent if model is None guard (line 127) creates a new db.TaskRun(...) instance, so whether scalar() returned a value or None, model is always orm_models.TaskRun after the else branch completes. After line 136, model is always non-None in both branches. All 3 errors are pyrefly-only, confirming this is a type narrowing limitation where pyrefly cannot track that both branches of the if/else guarantee a non-None assignment to model.
Attribution: The PR changes decompose_generator() and decompose_async_generator() in pyrefly/lib/alt/unwrap.rs to handle unions member-by-member, and adds normalize_async_generator_overloads() in pyrefly/lib/alt/function.rs. The union decomposition changes in unwrap.rs likely altered how unions are processed more broadly, which could have affected type narrowing for Union[TaskRun, None] in this unrelated code path. The refactoring of the union handling in decompose_generator (changing from imperative loop to functional collect/fold) may have introduced a subtle regression in how union types are narrowed through control flow.

✅ Improvement (3)

starlette (+4)

These are real type errors in starlette's test code. The functions are async generators (they use yield inside async def) but are annotated with -> ContentStream, which is a union type alias (AsyncGenerator[str | bytes, None] | Generator[str | bytes, None, None]). An async generator function should be annotated with AsyncGenerator (or AsyncIterator/AsyncIterable), not a union that also includes sync Generator. Pyright confirms all 4 errors. This is an improvement — pyrefly is now correctly catching a genuine annotation mismatch.
Attribution: The changes to decompose_async_generator() in pyrefly/lib/alt/unwrap.rs now handle Type::Union members individually. This more precise decomposition likely causes pyrefly to now properly validate that async generator return annotations must be AsyncGenerator-compatible rather than accepting union types that happen to contain AsyncGenerator as a member. The decompose_generator() changes in the same file follow the same pattern.

steam.py (-2)

Both removed errors were false positives. The code at lines 100-146 defines an overloaded async generator method entries(). The overload stubs (lines 101 and 110) and the implementation (line 114) all declare async def with return type AsyncGenerator[LeaderboardUser, None], and the implementation body uses yield. Pyrefly was incorrectly wrapping the overload stubs' return types as Coroutine[Unknown, Unknown, AsyncGenerator[LeaderboardUser]] before comparing them to the implementation's AsyncGenerator[LeaderboardUser], causing a spurious inconsistent-overload error. The PR correctly normalizes overload signatures when the implementation is an async generator, removing the false Coroutine wrapping.
Attribution: The fix is in pyrefly/lib/alt/function.rs where normalize_async_generator_overloads() was added. When the implementation signature is an async generator (detected via decompose_async_generator), the new code calls normalize_async_generator_overload_signature() on each overload stub. This method unwraps the Coroutine wrapper from the overload's return type when the inner type is an AsyncGenerator, so the overload return type matches the implementation's return type. The additional changes in pyrefly/lib/alt/unwrap.rs to decompose_async_generator() handle union types member-by-member, which is a related fix for union return types.

strawberry (+1)

The function subscribe at line 114-161 is an async generator function because it contains yield statements (lines 153, 156-158). For async generator functions in Python, the return type annotation should describe what the generator yields, typically AsyncIterator[ExecutionResult] or AsyncGenerator[ExecutionResult, None]. The annotation ExecutionResult | AsyncIterator[ExecutionResult] is incorrect because an async generator function cannot directly return an ExecutionResult value - it can only yield values and optionally return None. The bad-return error from pyrefly is a true positive: the return annotation includes ExecutionResult as a union member, but an async generator function's return type must be an async generator/iterator type, not a union containing non-generator types. This is a genuine misannotation in the source code. The function always acts as an async generator (yielding ExecutionResult values), so the correct annotation would be AsyncIterator[ExecutionResult]. The fact that other type checkers may not flag this doesn't make it a false positive - pyrefly is correctly identifying that the return type annotation is incompatible with the function being an async generator.
Attribution: The change to decompose_async_generator() in pyrefly/lib/alt/unwrap.rs added union member-by-member decomposition. When applied to ExecutionResult | AsyncIterator[ExecutionResult], it tries to decompose each member: AsyncIterator[ExecutionResult] succeeds but ExecutionResult fails, causing the overall decomposition to return None (via the ? on results?). This None result triggers the bad-return error. Previously, the function would have handled the whole type without union decomposition and likely taken a different code path.

Suggested fixes

Summary: The refactored decompose_async_generator() in unwrap.rs fails on union types containing non-async-generator members (like None), causing a false positive in freqtrade and likely affecting type narrowing in prefect; the optuna regression may also stem from broader union handling changes.

**1. In decompose_async_generator() in pyrefly/lib/alt/unwrap.rs, when handling Type::Union, filter out None/NoneType members before attempting to decompose each union member as an async generator. Currently, the code does members.iter().map(|member| self.decompose_async_generator(member)).collect() which returns None if ANY member fails decomposition. For return type annotations like AsyncIterator[RPC] | None, the None member is not an async generator and causes the entire decomposition to fail. The fix should skip None members (since async generators implicitly return None), e.g.:

rust
Type::Union(box Union { members, .. }) => {
let non_none_members: Vec<> = members.iter().filter(|m| !m.is_none()).collect();
if non_none_members.is_empty() {
return None;
}
let results: Option<Vec<>> = non_none_members
.iter()
.map(|member| self.decompose_async_generator(member))
.collect();
let (yield_tys, send_tys) = results?.into_iter().unzip();
Some((self.unions(yield_tys), self.unions(send_tys)))
}

This allows AsyncIterator[X] | None to be correctly recognized as an async generator return type, since the | None is semantically meaningless for async generators (they always return None implicitly).**

Files: pyrefly/lib/alt/unwrap.rs
Confidence: high
Affected projects: freqtrade
Fixes: bad-return
The freqtrade regression is directly caused by decompose_async_generator failing on AsyncIterator[RPC] | None because None cannot be decomposed as an async generator. The old code used is_subset_eq on the whole type which was more permissive. Filtering out None from the union before decomposition matches the semantic reality that async generators implicitly return None. This eliminates the 1 bad-return error in freqtrade. The same fix should also be applied to decompose_generator() for consistency.

2. In decompose_generator() in pyrefly/lib/alt/unwrap.rs, apply the same None-filtering fix as for decompose_async_generator(). The refactored code similarly fails if any union member is not a generator type. Filter out None/NoneType members before the .collect() call to handle Generator[X, Y, Z] | None patterns.

Files: pyrefly/lib/alt/unwrap.rs
Confidence: medium
Affected projects: optuna, prefect
Fixes: bad-argument, bad-return
The decompose_generator() function was refactored in the same way as decompose_async_generator() and has the same failure mode with union types containing None. While no direct regression is attributed to this specific function, it could be contributing to the optuna and prefect regressions through indirect type inference effects. The optuna regression (lost detection of str | None error) and prefect regression (3 new false positives on Union[TaskRun, None]) may be caused by changes in how unions flow through the type system after generator decomposition changes.

3. For the prefect regression specifically: investigate whether the union handling changes in decompose_generator()/decompose_async_generator() have side effects on general union type narrowing. The 3 new errors in prefect are about Union[orm_models.TaskRun, None] not being narrowed to orm_models.TaskRun after both branches of an if/else guarantee non-None assignment. This is a control-flow narrowing issue that may have been indirectly affected by the union decomposition refactoring. If the generator decomposition changes don't explain this, the regression may be a pre-existing narrowing limitation that was exposed by unrelated type resolution changes.

Files: pyrefly/lib/alt/unwrap.rs
Confidence: low
Affected projects: prefect
Fixes: bad-argument
The 3 prefect errors are all pyrefly-only and involve type narrowing of Optional types through control flow, which is not directly related to generator decomposition. The connection between the PR's changes and these errors is unclear. It's possible these are pre-existing issues that were coincidentally exposed, or that the union handling changes subtly affected type resolution order. More investigation is needed to confirm the causal link.

---

Was this helpful? React with 👍 or 👎

_{Classification by primer-classifier (6 LLM)}