[ty] Avoid unnecessary argument type expansion (#19999)

dhruvmanila · web-flow · commit d43a3d34dda1 · 2025-08-21T06:13:11.000Z
## Summary Part of: astral-sh/ty#868 This PR adds a heuristic to avoid argument type expansion if it's going to eventually lead to no matching overload. This is done by checking whether the non-expandable argument types are assignable to the corresponding annotated parameter type. If one of them is not assignable to all of the remaining overloads, then argument type expansion isn't going to help. ## Test Plan Add mdtest that would otherwise take a long time because of the number of arguments that it would need to expand (30).
diff --git a/crates/ty_python_semantic/resources/mdtest/call/overloads.md b/crates/ty_python_semantic/resources/mdtest/call/overloads.md
@@ -620,6 +620,217 @@ def _(ab: A | B, ac: A | C, cd: C | D):
     reveal_type(f(*(cd,)))  # revealed: Unknown
 ```
 
+### Optimization: Avoid argument type expansion
+
+Argument type expansion could lead to exponential growth of the number of argument lists that needs
+to be evaluated, so ty deploys some heuristics to prevent this from happening.
+
+Heuristic: If an argument type that cannot be expanded and cannot be assighned to any of the
+remaining overloads before argument type expansion, then even with argument type expansion, it won't
+lead to a successful evaluation of the call.
+
+`overloaded.pyi`:
+
+```pyi
+from typing import overload
+
+class A: ...
+class B: ...
+class C: ...
+
+@overload
+def f() -> None: ...
+@overload
+def f(**kwargs: int) -> C: ...
+@overload
+def f(x: A, /, **kwargs: int) -> A: ...
+@overload
+def f(x: B, /, **kwargs: int) -> B: ...
+
+class Foo:
+    @overload
+    def f(self) -> None: ...
+    @overload
+    def f(self, **kwargs: int) -> C: ...
+    @overload
+    def f(self, x: A, /, **kwargs: int) -> A: ...
+    @overload
+    def f(self, x: B, /, **kwargs: int) -> B: ...
+```
+
+```py
+from overloaded import A, B, C, Foo, f
+from typing_extensions import reveal_type
+
+def _(ab: A | B, a=1):
+    reveal_type(f(a1=a, a2=a, a3=a))  # revealed: C
+    reveal_type(f(A(), a1=a, a2=a, a3=a))  # revealed: A
+    reveal_type(f(B(), a1=a, a2=a, a3=a))  # revealed: B
+
+    # Here, the arity check filters out the first and second overload, type checking fails on the
+    # remaining overloads, so ty moves on to argument type expansion. But, the first argument (`C`)
+    # isn't assignable to any of the remaining overloads (3 and 4), so there's no point in expanding
+    # the other 30 arguments of type `Unknown | Literal[1]` which would result in allocating a
+    # vector containing 2**30 argument lists after expanding all of the arguments.
+    reveal_type(
+        # error: [no-matching-overload]
+        # revealed: Unknown
+        f(
+            C(),
+            a1=a,
+            a2=a,
+            a3=a,
+            a4=a,
+            a5=a,
+            a6=a,
+            a7=a,
+            a8=a,
+            a9=a,
+            a10=a,
+            a11=a,
+            a12=a,
+            a13=a,
+            a14=a,
+            a15=a,
+            a16=a,
+            a17=a,
+            a18=a,
+            a19=a,
+            a20=a,
+            a21=a,
+            a22=a,
+            a23=a,
+            a24=a,
+            a25=a,
+            a26=a,
+            a27=a,
+            a28=a,
+            a29=a,
+            a30=a,
+        )
+    )
+
+    # Here, the heuristics won't come into play because all arguments can be expanded but expanding
+    # the first argument resutls in a successful evaluation of the call, so there's no exponential
+    # growth of the number of argument lists.
+    reveal_type(
+        # revealed: A | B
+        f(
+            ab,
+            a1=a,
+            a2=a,
+            a3=a,
+            a4=a,
+            a5=a,
+            a6=a,
+            a7=a,
+            a8=a,
+            a9=a,
+            a10=a,
+            a11=a,
+            a12=a,
+            a13=a,
+            a14=a,
+            a15=a,
+            a16=a,
+            a17=a,
+            a18=a,
+            a19=a,
+            a20=a,
+            a21=a,
+            a22=a,
+            a23=a,
+            a24=a,
+            a25=a,
+            a26=a,
+            a27=a,
+            a28=a,
+            a29=a,
+            a30=a,
+        )
+    )
+
+def _(foo: Foo, ab: A | B, a=1):
+    reveal_type(foo.f(a1=a, a2=a, a3=a))  # revealed: C
+    reveal_type(foo.f(A(), a1=a, a2=a, a3=a))  # revealed: A
+    reveal_type(foo.f(B(), a1=a, a2=a, a3=a))  # revealed: B
+
+    reveal_type(
+        # error: [no-matching-overload]
+        # revealed: Unknown
+        foo.f(
+            C(),
+            a1=a,
+            a2=a,
+            a3=a,
+            a4=a,
+            a5=a,
+            a6=a,
+            a7=a,
+            a8=a,
+            a9=a,
+            a10=a,
+            a11=a,
+            a12=a,
+            a13=a,
+            a14=a,
+            a15=a,
+            a16=a,
+            a17=a,
+            a18=a,
+            a19=a,
+            a20=a,
+            a21=a,
+            a22=a,
+            a23=a,
+            a24=a,
+            a25=a,
+            a26=a,
+            a27=a,
+            a28=a,
+            a29=a,
+            a30=a,
+        )
+    )
+
+    reveal_type(
+        # revealed: A | B
+        foo.f(
+            ab,
+            a1=a,
+            a2=a,
+            a3=a,
+            a4=a,
+            a5=a,
+            a6=a,
+            a7=a,
+            a8=a,
+            a9=a,
+            a10=a,
+            a11=a,
+            a12=a,
+            a13=a,
+            a14=a,
+            a15=a,
+            a16=a,
+            a17=a,
+            a18=a,
+            a19=a,
+            a20=a,
+            a21=a,
+            a22=a,
+            a23=a,
+            a24=a,
+            a25=a,
+            a26=a,
+            a27=a,
+            a28=a,
+            a29=a,
+            a30=a,
+        )
+    )
+```
+
 ## Filtering based on `Any` / `Unknown`
 
 This is the step 5 of the overload call evaluation algorithm which specifies that:
diff --git a/crates/ty_python_semantic/src/types/call/arguments.rs b/crates/ty_python_semantic/src/types/call/arguments.rs
@@ -5,7 +5,7 @@ use ruff_python_ast as ast;
 
 use crate::Db;
 use crate::types::KnownClass;
-use crate::types::enums::enum_member_literals;
+use crate::types::enums::{enum_member_literals, enum_metadata};
 use crate::types::tuple::{Tuple, TupleLength, TupleType};
 
 use super::Type;
@@ -208,10 +208,32 @@ impl<'a, 'db> FromIterator<(Argument<'a>, Option<Type<'db>>)> for CallArguments<
     }
 }
 
+/// Returns `true` if the type can be expanded into its subtypes.
+///
+/// In other words, it returns `true` if [`expand_type`] returns [`Some`] for the given type.
+pub(crate) fn is_expandable_type<'db>(db: &'db dyn Db, ty: Type<'db>) -> bool {
+    match ty {
+        Type::NominalInstance(instance) => {
+            let class = instance.class(db);
+            class.is_known(db, KnownClass::Bool)
+                || instance.tuple_spec(db).is_some_and(|spec| match &*spec {
+                    Tuple::Fixed(fixed_length_tuple) => fixed_length_tuple
+                        .all_elements()
+                        .any(|element| is_expandable_type(db, *element)),
+                    Tuple::Variable(_) => false,
+                })
+                || enum_metadata(db, class.class_literal(db).0).is_some()
+        }
+        Type::Union(_) => true,
+        _ => false,
+    }
+}
+
 /// Expands a type into its possible subtypes, if applicable.
 ///
 /// Returns [`None`] if the type cannot be expanded.
 fn expand_type<'db>(db: &'db dyn Db, ty: Type<'db>) -> Option<Vec<Type<'db>>> {
+    // NOTE: Update `is_expandable_type` if this logic changes accordingly.
     match ty {
         Type::NominalInstance(instance) => {
             let class = instance.class(db);
diff --git a/crates/ty_python_semantic/src/types/call/bind.rs b/crates/ty_python_semantic/src/types/call/bind.rs
@@ -16,6 +16,7 @@ use crate::Program;
 use crate::db::Db;
 use crate::dunder_all::dunder_all_names;
 use crate::place::{Boundness, Place};
+use crate::types::call::arguments::is_expandable_type;
 use crate::types::diagnostic::{
     CALL_NON_CALLABLE, CONFLICTING_ARGUMENT_FORMS, INVALID_ARGUMENT_TYPE, MISSING_ARGUMENT,
     NO_MATCHING_OVERLOAD, PARAMETER_ALREADY_ASSIGNED, TOO_MANY_POSITIONAL_ARGUMENTS,
@@ -1337,6 +1338,48 @@ impl<'db> CallableBinding<'db> {
         // for evaluating the expanded argument lists.
         snapshotter.restore(self, pre_evaluation_snapshot);
 
+        // At this point, there's at least one argument that can be expanded.
+        //
+        // This heuristic tries to detect if there's any need to perform argument type expansion or
+        // not by checking whether there are any non-expandable argument type that cannot be
+        // assigned to any of the remaining overloads.
+        //
+        // This heuristic needs to be applied after restoring the bindings state to the one before
+        // type checking as argument type expansion would evaluate it from that point on.
+        for (argument_index, (argument, argument_type)) in argument_types.iter().enumerate() {
+            // TODO: Remove `Keywords` once `**kwargs` support is added
+            if matches!(argument, Argument::Synthetic | Argument::Keywords) {
+                continue;
+            }
+            let Some(argument_type) = argument_type else {
+                continue;
+            };
+            if is_expandable_type(db, argument_type) {
+                continue;
+            }
+            let mut is_argument_assignable_to_any_overload = false;
+            'overload: for (_, overload) in self.matching_overloads() {
+                for parameter_index in &overload.argument_matches[argument_index].parameters {
+                    let parameter_type = overload.signature.parameters()[*parameter_index]
+                        .annotated_type()
+                        .unwrap_or(Type::unknown());
+                    if argument_type.is_assignable_to(db, parameter_type) {
+                        is_argument_assignable_to_any_overload = true;
+                        break 'overload;
+                    }
+                }
+            }
+            if !is_argument_assignable_to_any_overload {
+                tracing::debug!(
+                    "Argument at {argument_index} (`{}`) is not assignable to any of the \
+                    remaining overloads, skipping argument type expansion",
+                    argument_type.display(db)
+                );
+                snapshotter.restore(self, post_evaluation_snapshot);
+                return;
+            }
+        }
+
         for expanded_argument_lists in expansions {
             // This is the merged state of the bindings after evaluating all of the expanded
             // argument lists. This will be the final state to restore the bindings to if all of
