fix

Author: asukaminato0721

pyrefly/lib/alt/function.rs

@@ -375,6 +375,11 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                         defs,
                         errors,
                     );
+                    let sigs = if let [impl_sig] = def.ty.callable_signatures().as_slice() {
+                        self.normalize_async_generator_overloads(sigs, impl_sig)
+                    } else {
+                        sigs
+                    };
                     self.check_signature_consistency(&sigs, &def, errors);
                     Type::Overload(Overload {
                         signatures: sigs.mapped(|(_, sig)| sig),
@@ -1704,6 +1709,49 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
         metadata
     }
 
+    fn normalize_async_generator_overload_signature(&self, sig: &Callable) -> Callable {
+        if let Some((_, _, ret)) = self.unwrap_coroutine(&sig.ret)
+            && self.decompose_async_generator(&ret).is_some()
+        {
+            let mut sig = sig.clone();
+            sig.ret = ret;
+            sig
+        } else {
+            sig.clone()
+        }
+    }
+
+    fn normalize_async_generator_overloads(
+        &self,
+        overloads: Vec1<(TextRange, OverloadType)>,
+        impl_sig: &Callable,
+    ) -> Vec1<(TextRange, OverloadType)> {
+        if self.decompose_async_generator(&impl_sig.ret).is_none() {
+            return overloads;
+        }
+        overloads.mapped(|(range, overload)| {
+            (
+                range,
+                match overload {
+                    OverloadType::Function(func) => OverloadType::Function(Function {
+                        signature: self
+                            .normalize_async_generator_overload_signature(&func.signature),
+                        metadata: func.metadata,
+                    }),
+                    OverloadType::Forall(forall) => OverloadType::Forall(Forall {
+                        tparams: forall.tparams,
+                        body: Function {
+                            signature: self.normalize_async_generator_overload_signature(
+                                &forall.body.signature,
+                            ),
+                            metadata: forall.body.metadata,
+                        },
+                    }),
+                },
+            )
+        })
+    }
+
     /// Substitute each type parameter with its upper bound, for overload consistency checking.
     fn subst_function(&self, tparams: &TParams, func: Function) -> Function {
         let mp = tparams

pyrefly/lib/alt/unwrap.rs

@@ -375,16 +375,20 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
 
     pub fn decompose_generator(&self, ty: &Type) -> Option<(Type, Type, Type)> {
         match ty {
-            Type::Union(u) => {
-                let mut yield_tys = Vec::new();
-                let mut send_tys = Vec::new();
-                let mut return_tys = Vec::new();
-                for member in &u.members {
-                    let (y, s, r) = self.decompose_generator(member)?;
-                    yield_tys.push(y);
-                    send_tys.push(s);
-                    return_tys.push(r);
-                }
+            Type::Union(box Union { members, .. }) => {
+                let results: Option<Vec<_>> = members
+                    .iter()
+                    .map(|member| self.decompose_generator(member))
+                    .collect();
+                let (yield_tys, send_tys, return_tys) = results?.into_iter().fold(
+                    (Vec::new(), Vec::new(), Vec::new()),
+                    |mut acc, (yield_ty, send_ty, return_ty)| {
+                        acc.0.push(yield_ty);
+                        acc.1.push(send_ty);
+                        acc.2.push(return_ty);
+                        acc
+                    },
+                );
                 Some((
                     self.unions(yield_tys),
                     self.unions(send_tys),
@@ -417,22 +421,34 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
     }
 
     pub fn decompose_async_generator(&self, ty: &Type) -> Option<(Type, Type)> {
-        let yield_ty = self.fresh_var();
-        let send_ty = self.fresh_var();
-        let async_generator_ty = self.heap.mk_class_type(
-            self.stdlib
-                .async_generator(yield_ty.to_type(self.heap), send_ty.to_type(self.heap)),
-        );
-        if self.is_subset_eq(&async_generator_ty, ty) {
-            let yield_ty: Type = self.resolve_var_opt(ty, yield_ty)?;
-            let send_ty = self
-                .resolve_var_opt(ty, send_ty)
-                .unwrap_or_else(|| self.heap.mk_none());
-            Some((yield_ty, send_ty))
-        } else if ty.is_any() {
-            Some((self.heap.mk_any_explicit(), self.heap.mk_any_explicit()))
-        } else {
-            None
+        match ty {
+            Type::Union(box Union { members, .. }) => {
+                let results: Option<Vec<_>> = 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)))
+            }
+            _ => {
+                let yield_ty = self.fresh_var();
+                let send_ty = self.fresh_var();
+                let async_generator_ty = self.heap.mk_class_type(
+                    self.stdlib
+                        .async_generator(yield_ty.to_type(self.heap), send_ty.to_type(self.heap)),
+                );
+                if self.is_subset_eq(&async_generator_ty, ty) {
+                    let yield_ty: Type = self.resolve_var_opt(ty, yield_ty)?;
+                    let send_ty = self
+                        .resolve_var_opt(ty, send_ty)
+                        .unwrap_or_else(|| self.heap.mk_none());
+                    Some((yield_ty, send_ty))
+                } else if ty.is_any() {
+                    Some((self.heap.mk_any_explicit(), self.heap.mk_any_explicit()))
+                } else {
+                    None
+                }
+            }
         }
     }
 

pyrefly/lib/test/overload.rs

@@ -834,6 +834,34 @@ def f[T](x: T) -> T:
     "#,
 );
 
+testcase!(
+    test_overloaded_async_iterator_impl,
+    r#"
+from collections.abc import AsyncIterator
+from typing import Literal, assert_type, overload
+
+class Patch:
+    pass
+
+class FullState(Patch):
+    pass
+
+@overload
+async def astream(*, diff: Literal[True] = ...) -> AsyncIterator[Patch]: ...
+@overload
+async def astream(*, diff: Literal[False]) -> AsyncIterator[FullState]: ...
+async def astream(*, diff: bool = True) -> AsyncIterator[Patch] | AsyncIterator[FullState]:
+    if diff:
+        yield Patch()
+    else:
+        yield FullState()
+
+assert_type(astream(), AsyncIterator[Patch])
+assert_type(astream(diff=True), AsyncIterator[Patch])
+assert_type(astream(diff=False), AsyncIterator[FullState])
+    "#,
+);
+
 testcase!(
     test_param_consistency,
     r#"