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

## Summary

Part of: https://github.com/astral-sh/ty/issues/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).
This commit is contained in:
Dhruv Manilawala 2025-08-21 11:43:11 +05:30 committed by GitHub
parent 99111961c0
commit d43a3d34dd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 277 additions and 1 deletions

View file

@ -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:

View file

@ -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);

View file

@ -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