Skip to content

[ty] Avoid unnecessary argument type expansion #19999

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Aug 21, 2025
Merged

[ty] Avoid unnecessary argument type expansion #19999

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
89f1298
[ty] Avoid unnecessary argument type expansion
dhruvmanila Aug 20, 2025
d52ac51
Add another test
dhruvmanila Aug 20, 2025
e89ae92
Fix pre-commit
dhruvmanila Aug 20, 2025
2e00f75
Avoid `**kwargs`
dhruvmanila Aug 20, 2025
5345af0
Add tests containing `self`
dhruvmanila Aug 20, 2025
1c99150
Address review comments
dhruvmanila Aug 21, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
211 changes: 211 additions & 0 deletions crates/ty_python_semantic/resources/mdtest/call/overloads.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Comment on lines +675 to +679
Copy link
Member Author

@dhruvmanila dhruvmanila Aug 20, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Considering the same example, but pass in A() as the first argument and use a=None instead of a=1 for the parameter above, this would still continue with the expansion. The reason being that A() matches parameter of at least one overload which means we cannot skip argument type expansion.

I'm trying to think through what kind of heuristic to have that would avoid the expansion in this case as well. This would still lead to no-matching-overload because even though Unknown is assignable to int, the spec says that all argument lists resulting from an expansion should evaluate successfully and there's no such expansion as expanding Unknown | None leads to argument lists where at least one would have None which isn't assignable to int.

Regardless, I think it'd be best to set a higher limit on the expansion.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I think there will be cases where we can't avoid actually trying all the expansions and see if any of them work -- best we can do there is optimize the expansions as best we can (e.g. iterator instead of eagerly materialized?) and then set a reasonable limit.

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:
Expand Down
24 changes: 23 additions & 1 deletion crates/ty_python_semantic/src/types/call/arguments.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down
43 changes: 43 additions & 0 deletions crates/ty_python_semantic/src/types/call/bind.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -1327,6 +1328,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.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would it make sense to first explicitly check if there are any expandable arguments (using the new method you added), then do this check, then actually expand the arguments? This way even if the heuristic applies, we've already generated the expansion, just haven't used the expansions yet.

But maybe generating the expansions is very cheap relative to actually trying the call with new argument types, so this doesn't matter?

And I guess we'd slow down the fast path slightly if we first check all arguments for expandability, then separately expand them?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I don't think this would really matter in practice mainly because there are only a few differences between expand_type and is_type_expandable which are (1) collecting the enum members from the metadata (both method generates the metadata) (2) performing multi-cartesian product of tuple elements and (3) allocation. I guess it wouldn't hurt to avoid allocating when it's easy to do so. Let me try it, I think I'll make this change.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah right, so the main reason I did it at this location is so that the bindings state is the one before the type checking step, so that we only skip the overloads that have been filtered by the arity check.

So, we'd either have to pay the cost of either

  • Allocation for the first expansion, taking the bindings snapshot. We'd skip the snapshot if there are no expansion in this case.
  • Always take the binding snapshot, check whether expansion needs to happen using the logic in this PR, skip allocating for the first expansion if there's no need

I think going with the latter sounds reasonable i.e., instead of always allocating the first expansion, we'd always allocate the bindings snapshot.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, but then there's complication on restoring the snapshots at the correct places. I think I'll leave this as is for now.

//
// 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
Expand Down
Loading