use argument type for target overload when filtering overloads with gradual types

Author: ibraheemdev

crates/ty_python_semantic/resources/mdtest/bidirectional.md

@@ -319,6 +319,17 @@ def list_or_set2[T, U](x: T, y: U) -> list[T] | set[U]:
 # error: [no-matching-overload]
 x8 = f6(reveal_type(list_or_set2(1, 1)))  # revealed: list[int] | set[int]
 reveal_type(x8)  # revealed: Unknown
+
+@overload
+def f7(y: list[int | str]) -> list[int | str]: ...
+@overload
+def f7[T](y: list[T]) -> list[T]: ...
+def f7(y: object) -> object:
+    raise NotImplementedError
+
+# TODO: We should reveal `list[int | str]` here.
+x9 = f7(reveal_type(["Sheet1"]))  # revealed: list[str]
+reveal_type(x9)  # revealed: list[int | str]
 ```
 
 ## Class constructor parameters

crates/ty_python_semantic/src/types/call/arguments.rs

@@ -53,14 +53,14 @@ pub(crate) struct CallArguments<'a, 'db> {
 /// with type context across multiple bindings.
 #[derive(Clone, Debug, Default)]
 pub(crate) struct CallArgumentTypes<'db> {
-    fallback_ty: Option<Type<'db>>,
+    fallback_type: Option<Type<'db>>,
     types: FxHashMap<Type<'db>, Type<'db>>,
 }
 
 impl<'db> CallArgumentTypes<'db> {
     pub(crate) fn new(fallback_ty: Option<Type<'db>>) -> Self {
         Self {
-            fallback_ty,
+            fallback_type: fallback_ty,
             types: FxHashMap::default(),
         }
     }
@@ -78,7 +78,7 @@ impl<'db> CallArgumentTypes<'db> {
             return Some(*exact_ty);
         }
 
-        self.fallback_ty
+        self.fallback_type
     }
 
     /// Returns the type of this argument when inferred against the provided declared type.
@@ -93,7 +93,7 @@ impl<'db> CallArgumentTypes<'db> {
     /// Insert the type of this argument when inferred with the provided type context.
     pub(crate) fn insert(&mut self, tcx: impl Into<TypeContext<'db>>, ty: Type<'db>) {
         match tcx.into().annotation {
-            None => self.fallback_ty = Some(ty),
+            None => self.fallback_type = Some(ty),
             Some(tcx) => {
                 self.types.insert(tcx, ty);
             }
@@ -104,7 +104,7 @@ impl<'db> CallArgumentTypes<'db> {
         self.types
             .iter()
             .map(|(tcx, ty)| (TypeContext::new(Some(*tcx)), *ty))
-            .chain(self.fallback_ty.map(|ty| (TypeContext::default(), ty)))
+            .chain(self.fallback_type.map(|ty| (TypeContext::default(), ty)))
     }
 }
 

crates/ty_python_semantic/src/types/call/bind.rs

@@ -2590,71 +2590,6 @@ impl<'db> CallableBinding<'db> {
             }
         }
 
-        let mut union_argument_type_builders = std::iter::repeat_with(|| UnionBuilder::new(db))
-            .take(max_parameter_count)
-            .collect::<Vec<_>>();
-
-        // The following loop is trying to construct a tuple of argument types that correspond to
-        // the participating parameter indexes. Considering the following example:
-        //
-        // ```python
-        // @overload
-        // def f(x: Literal[1], y: Literal[2]) -> tuple[int, int]: ...
-        // @overload
-        // def f(*args: Any) -> tuple[Any, ...]: ...
-        //
-        // f(1, 2)
-        // ```
-        //
-        // Here, only the first parameter participates in the filtering process because only one
-        // overload has the second parameter. So, while going through the argument types, the
-        // second argument needs to be skipped but for the second overload both arguments map to
-        // the first parameter and that parameter is considered for the filtering process. This
-        // flag is to handle that special case of many-to-one mapping from arguments to parameters.
-        let mut variadic_parameter_handled = false;
-
-        for (argument_index, argument_types) in arguments.types().iter().enumerate() {
-            if variadic_parameter_handled {
-                continue;
-            }
-            for overload_index in matching_overload_indexes {
-                let overload = &self.overloads[*overload_index];
-                for (parameter_index, variadic_argument_type) in
-                    overload.argument_matches[argument_index].iter()
-                {
-                    let parameter = &overload.signature.parameters()[parameter_index];
-                    if parameter.is_variadic() {
-                        variadic_parameter_handled = true;
-                    }
-                    if !participating_parameter_indexes.contains(&parameter_index) {
-                        continue;
-                    }
-                    let argument_type =
-                        argument_types.get_for_declared_type(parameter.annotated_type());
-                    union_argument_type_builders[parameter_index].add_in_place(
-                        variadic_argument_type
-                            .unwrap_or(argument_type)
-                            .top_materialization(db),
-                    );
-                }
-            }
-        }
-
-        // These only contain the top materialized argument types for the corresponding
-        // participating parameter indexes.
-        let top_materialized_argument_type = Type::heterogeneous_tuple(
-            db,
-            union_argument_type_builders
-                .into_iter()
-                .filter_map(|builder| {
-                    if builder.is_empty() {
-                        None
-                    } else {
-                        Some(builder.build())
-                    }
-                }),
-        );
-
         // A flag to indicate whether we've found the overload that makes the remaining overloads
         // unmatched for the given argument types.
         let mut filter_remaining_overloads = false;
@@ -2665,6 +2600,84 @@ impl<'db> CallableBinding<'db> {
                 continue;
             }
 
+            let mut union_argument_type_builders = std::iter::repeat_with(|| UnionBuilder::new(db))
+                .take(max_parameter_count)
+                .collect::<Vec<_>>();
+
+            // The following loop is trying to construct a tuple of argument types that correspond to
+            // the participating parameter indexes. Considering the following example:
+            //
+            // ```python
+            // @overload
+            // def f(x: Literal[1], y: Literal[2]) -> tuple[int, int]: ...
+            // @overload
+            // def f(*args: Any) -> tuple[Any, ...]: ...
+            //
+            // f(1, 2)
+            // ```
+            //
+            // Here, only the first parameter participates in the filtering process because only one
+            // overload has the second parameter. So, while going through the argument types, the
+            // second argument needs to be skipped but for the second overload both arguments map to
+            // the first parameter and that parameter is considered for the filtering process. This
+            // flag is to handle that special case of many-to-one mapping from arguments to parameters.
+            let mut variadic_parameter_handled = false;
+
+            for (argument_index, argument_types) in arguments.types().iter().enumerate() {
+                if variadic_parameter_handled {
+                    continue;
+                }
+
+                // Get the argument type as inferred against the target overload.
+                let current_overload = &self.overloads[*current_index];
+                let argument_type =
+                    match *current_overload.argument_matches[argument_index].parameters {
+                        [parameter_index] => {
+                            let declared_type = current_overload.signature.parameters()
+                                [parameter_index]
+                                .annotated_type();
+                            argument_types.get_for_declared_type(declared_type)
+                        }
+                        // Splatted arguments are inferred without type context.
+                        _ => argument_types.get_default().unwrap_or(Type::unknown()),
+                    };
+
+                for overload_index in matching_overload_indexes {
+                    let overload = &self.overloads[*overload_index];
+                    for (parameter_index, variadic_argument_type) in
+                        overload.argument_matches[argument_index].iter()
+                    {
+                        let parameter = &overload.signature.parameters()[parameter_index];
+                        if parameter.is_variadic() {
+                            variadic_parameter_handled = true;
+                        }
+                        if !participating_parameter_indexes.contains(&parameter_index) {
+                            continue;
+                        }
+                        union_argument_type_builders[parameter_index].add_in_place(
+                            variadic_argument_type
+                                .unwrap_or(argument_type)
+                                .top_materialization(db),
+                        );
+                    }
+                }
+            }
+
+            // These only contain the top materialized argument types for the corresponding
+            // participating parameter indexes.
+            let top_materialized_argument_type = Type::heterogeneous_tuple(
+                db,
+                union_argument_type_builders
+                    .into_iter()
+                    .filter_map(|builder| {
+                        if builder.is_empty() {
+                            None
+                        } else {
+                            Some(builder.build())
+                        }
+                    }),
+            );
+
             let mut union_parameter_types = std::iter::repeat_with(|| UnionBuilder::new(db))
                 .take(max_parameter_count)
                 .collect::<Vec<_>>();