[ty] Retry parameter matching for argument type expansion

crates/ty_python_semantic/resources/mdtest/call/function.md

@@ -613,56 +613,6 @@ def _(args: str) -> None:
     takes_at_least_two_positional_only(*args)
 ```
 
-### Argument expansion regression
-
-This is a regression that was highlighted by the ecosystem check, which shows that we might need to
-rethink how we perform argument expansion during overload resolution. In particular, we might need
-to retry both `match_parameters` *and* `check_types` for each expansion. Currently we only retry
-`check_types`.
-
-The issue is that argument expansion might produce a splatted value with a different arity than what
-we originally inferred for the unexpanded value, and that in turn can affect which parameters the
-splatted value is matched with.
-
-The first example correctly produces an error. The `tuple[int, str]` union element has a precise
-arity of two, and so parameter matching chooses the first overload. The second element of the tuple
-does not match the second parameter type, which yielding an `invalid-argument-type` error.
-
-The third example should produce the same error. However, because we have a union, we do not see the
-precise arity of each union element during parameter matching. Instead, we infer an arity of "zero
-or more" for the union as a whole, and use that less precise arity when matching parameters. We
-therefore consider the second overload to still be a potential candidate for the `tuple[int, str]`
-union element. During type checking, we have to force the arity of each union element to match the
-inferred arity of the union as a whole (turning `tuple[int, str]` into `tuple[int | str, ...]`).
-That less precise tuple type-checks successfully against the second overload, making us incorrectly
-think that `tuple[int, str]` is a valid splatted call.
-
-If we update argument expansion to retry parameter matching with the precise arity of each union
-element, we will correctly rule out the second overload for `tuple[int, str]`, just like we do when
-splatting that tuple directly (instead of as part of a union).
-
-```py
-from typing import overload
-
-@overload
-def f(x: int, y: int) -> None: ...
-@overload
-def f(x: int, y: str, z: int) -> None: ...
-def f(*args): ...
-
-# Test all of the above with a number of different splatted argument types
-
-def _(t: tuple[int, str]) -> None:
-    f(*t)  # error: [invalid-argument-type]
-
-def _(t: tuple[int, str, int]) -> None:
-    f(*t)
-
-def _(t: tuple[int, str] | tuple[int, str, int]) -> None:
-    # TODO: error: [invalid-argument-type]
-    f(*t)
-```
-
 ## Wrong argument type
 
 ### Positional argument, positional-or-keyword parameter

crates/ty_python_semantic/resources/mdtest/call/overloads.md

@@ -889,6 +889,48 @@ def _(a: int | None):
     )
 ```
 
+### Retry from parameter matching
+
+As per the spec, the argument type expansion should retry evaluating the expanded argument list from
+the type checking step. However, that creates an issue when variadic arguments are involved because
+if a variadic argument is a union type, it could be expanded to have different arities. So, ty
+retries it from the start which includes parameter matching as well.
+
+`overloaded.pyi`:
+
+```pyi
+from typing import overload
+
+@overload
+def f(x: int, y: int) -> None: ...
+@overload
+def f(x: int, y: str, z: int) -> None: ...
+```
+
+```py
+from overloaded import f
+
+# Test all of the above with a number of different splatted argument types
+
+def _(t: tuple[int, str]) -> None:
+    # This correctly produces an error because the first element of the union has a precise arity of
+    # 2, which matches the first overload, but the second element of the tuple doesn't match the
+    # second parameter type, yielding an `invalid-argument-type` error.
+    f(*t)  # error: [invalid-argument-type]
+
+def _(t: tuple[int, str, int]) -> None:
+    # This correctly produces no error because the first element of the union has a precise arity of
+    # 3, which matches the second overload.
+    f(*t)
+
+def _(t: tuple[int, str] | tuple[int, str, int]) -> None:
+    # This produces an error because the expansion produces two argument lists: `[*tuple[int, str]]`
+    # and `[*tuple[int, str, int]]`. The first list produces produces a type checking error as
+    # described in the first example, while the second list matches the second overload. And,
+    # because not all of the expanded argument list evaluates successfully, we produce an error.
+    f(*t)  # error: [no-matching-overload]
+```
+
 ## Filtering based on `Any` / `Unknown`
 
 This is the step 5 of the overload call evaluation algorithm which specifies that:

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

@@ -202,10 +202,25 @@ impl<'a, 'db> CallArguments<'a, 'db> {
                     for subtype in &expanded_types {
                         let mut new_expanded_types = pre_expanded_types.to_vec();
                         new_expanded_types[index] = Some(*subtype);
-                        expanded_arguments.push(CallArguments::new(
-                            self.arguments.clone(),
-                            new_expanded_types,
-                        ));
+
+                        // Update the arguments list to handle variadic argument expansion
+                        let mut new_arguments = self.arguments.clone();
+                        if let Argument::Variadic(_) = self.arguments[index] {
+                            // If the argument corresponding to this type is variadic, we need to
+                            // update the tuple length because expanding could change the length.
+                            // For example, in `tuple[int] | tuple[int, int]`, the length of the
+                            // first type is 1, while the length of the second type is 2.
+                            if let Some(expanded_type) = new_expanded_types[index] {
+                                let length = expanded_type
+                                    .try_iterate(db)
+                                    .map(|tuple| tuple.len())
+                                    .unwrap_or(TupleLength::unknown());
+                                new_arguments[index] = Argument::Variadic(length);
+                            }
+                        }
+
+                        expanded_arguments
+                            .push(CallArguments::new(new_arguments, new_expanded_types));
                     }
                 }