diff --git a/crates/ty_python_semantic/resources/mdtest/call/overloads.md b/crates/ty_python_semantic/resources/mdtest/call/overloads.md index ca447b6063..d6fbf08b3a 100644 --- 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 index b200f3fc8a..648b0cdab9 100644 --- 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>)> 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>> { + // 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 index d904e4c9ab..da685d68ed 100644 --- 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