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:

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]

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

❌ 1 regression(s) | ✅ 1 improvement(s) | 2 project(s) total | +3, -2 errors

1 regression(s) across prefect. error kinds: bad-return, missing-attribute. caused by decompose_async_generator(). 1 improvement(s) across steam.py.

Project	Verdict	Changes	Error Kinds	Root Cause
steam.py	✅ Improvement	-2	inconsistent-overload	normalize_async_generator_overloads()
prefect	❌ Regression	+3	bad-return, missing-attribute	decompose_async_generator()

Detailed analysis

❌ Regression (1)

prefect (+3)

All three errors stem from pyrefly failing to narrow model: Union[orm_models.TaskRun, None] to orm_models.TaskRun after a conditional assignment pattern where the if model is None branch reassigns model to a non-None value. This is a false positive — the code correctly ensures model is never None at line 138. The missing-attribute errors on .created and .id are cascading from the same root cause (pyrefly thinks model could be None). The bad-return error similarly follows from the unnarrowed type. All 3 errors are pyrefly-only, confirming this is a regression in pyrefly's type narrowing.
Attribution: The PR changes 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. These changes to union handling in the unwrap/function modules likely affected how pyrefly resolves types through complex control flow, causing it to lose narrowing information for the model variable after the if/else branches. The union decomposition changes may have inadvertently affected how union types are tracked through conditional assignments.

✅ Improvement (1)

steam.py (-2)

These are false positives caused by a nuanced interaction between async def and async generator functions. The overload stubs (lines 101 and 110) are declared as async def entries(...) -> AsyncGenerator[LeaderboardUser, None]: .... Since these stubs contain no yield statement (just ...), a type checker applying standard rules would treat them as regular coroutine functions, inferring their actual return type as Coroutine[Unknown, Unknown, AsyncGenerator[LeaderboardUser]] — an async function that returns an AsyncGenerator when awaited.

The implementation (line 114) is also declared as async def entries(...) -> AsyncGenerator[LeaderboardUser, None], but its body contains yield statements (line 139), making it an async generator function. For async generator functions, the return type annotation is interpreted directly as AsyncGenerator[LeaderboardUser, None] without Coroutine wrapping.

This creates the type mismatch: pyrefly sees the overload stubs as having return type Coroutine[Unknown, Unknown, AsyncGenerator[LeaderboardUser]] while the implementation has return type AsyncGenerator[LeaderboardUser, None]. However, this is a false positive because the overload stubs are intended to describe the same async generator behavior as the implementation — they just can't contain yield in a stub. The PR correctly handles this case by recognizing that when comparing overload return types against an async generator implementation, the AsyncGenerator return annotation on async def stubs should not be wrapped in Coroutine.

Attribution: The fix is in pyrefly/lib/alt/function.rs where normalize_async_generator_overloads() and normalize_async_generator_overload_signature() were added. When the implementation signature is an async generator, the new code unwraps the Coroutine[..., AsyncGenerator[...]] wrapper from each overload stub's return type back to just AsyncGenerator[...] before performing the consistency check. This directly eliminates the false positive inconsistent-overload errors.

Suggested fixes

Summary: The decompose_async_generator() change in unwrap.rs adds a Type::Union match arm that short-circuits before calling is_subset_eq/fresh_var for non-async-generator unions, potentially affecting solver constraint state and causing narrowing regressions in prefect.

1. In decompose_async_generator() in pyrefly/lib/alt/unwrap.rs, the new Type::Union match arm should only recurse into union members when at least one member could plausibly be an AsyncGenerator. For unions like Union[TaskRun, None] where no member is an AsyncGenerator, the function should fall through to the original default branch (which returns None quickly). Add a guard: before recursing, check if any union member is a class type that could be an AsyncGenerator (or is Any). If none are, fall through to the default _ arm. Alternatively, ensure the _ arm is reached for non-async-generator unions by checking if u.members.iter().any(|m| self.could_be_async_generator(m)) before the recursive decomposition.

Files: pyrefly/lib/alt/unwrap.rs
Confidence: medium
Affected projects: prefect
Fixes: missing-attribute, bad-return
The old code for Union[TaskRun, None] would hit the default branch and call is_subset_eq (which would fail and return None). The new code hits the Type::Union branch, recursively calls decompose_async_generator on each member (TaskRun returns None, causing collect() to return None). While both return None, the old path called is_subset_eq with fresh type variables against the union, potentially adding solver constraints that aided narrowing elsewhere. The new path skips this entirely. Fixing this to only use the Union branch for actual async generator unions would restore the old behavior for irrelevant union types, eliminating 3 pyrefly-only errors in prefect.

---

Was this helpful? React with 👍 or 👎

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