[ty] Use annotated parameters as type context (#20635)

## Summary

Use the type annotation of function parameters as bidirectional type
context when inferring the argument expression. For example, the
following example now type-checks:

```py
class TD(TypedDict):
    x: int

def f(_: TD): ...

f({ "x": 1 })
```

Part of astral-sh/ty#168.

Author: ibraheemdev

crates/ruff_benchmark/benches/ty_walltime.rs

@@ -117,7 +117,7 @@ static COLOUR_SCIENCE: std::sync::LazyLock<Benchmark<'static>> = std::sync::Lazy
             max_dep_date: "2025-06-17",
             python_version: PythonVersion::PY310,
         },
-        500,
+        600,
     )
 });
 

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

@@ -1662,3 +1662,67 @@ def _(arg: tuple[A | B, Any]):
     reveal_type(f(arg))  # revealed: Unknown
     reveal_type(f(*(arg,)))  # revealed: Unknown
 ```
+
+## Bidirectional Type Inference
+
+```toml
+[environment]
+python-version = "3.12"
+```
+
+Type inference accounts for parameter type annotations across all overloads.
+
+```py
+from typing import TypedDict, overload
+
+class T(TypedDict):
+    x: int
+
+@overload
+def f(a: list[T], b: int) -> int: ...
+@overload
+def f(a: list[dict[str, int]], b: str) -> str: ...
+def f(a: list[dict[str, int]] | list[T], b: int | str) -> int | str:
+    return 1
+
+def int_or_str() -> int | str:
+    return 1
+
+x = f([{"x": 1}], int_or_str())
+reveal_type(x)  # revealed: int | str
+
+# TODO: error: [no-matching-overload] "No overload of function `f` matches arguments"
+# we currently incorrectly consider `list[dict[str, int]]` a subtype of `list[T]`
+f([{"y": 1}], int_or_str())
+```
+
+Non-matching overloads do not produce diagnostics:
+
+```py
+from typing import TypedDict, overload
+
+class T(TypedDict):
+    x: int
+
+@overload
+def f(a: T, b: int) -> int: ...
+@overload
+def f(a: dict[str, int], b: str) -> str: ...
+def f(a: T | dict[str, int], b: int | str) -> int | str:
+    return 1
+
+x = f({"y": 1}, "a")
+reveal_type(x)  # revealed: str
+```
+
+```py
+from typing import SupportsRound, overload
+
+@overload
+def takes_str_or_float(x: str): ...
+@overload
+def takes_str_or_float(x: float): ...
+def takes_str_or_float(x: float | str): ...
+
+takes_str_or_float(round(1.0))
+```

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

@@ -251,3 +251,59 @@ from ty_extensions import Intersection, Not
 def _(x: Union[Intersection[Any, Not[int]], Intersection[Any, Not[int]]]):
     reveal_type(x)  # revealed: Any & ~int
 ```
+
+## Bidirectional Type Inference
+
+```toml
+[environment]
+python-version = "3.12"
+```
+
+Type inference accounts for parameter type annotations across all signatures in a union.
+
+```py
+from typing import TypedDict, overload
+
+class T(TypedDict):
+    x: int
+
+def _(flag: bool):
+    if flag:
+        def f(x: T) -> int:
+            return 1
+    else:
+        def f(x: dict[str, int]) -> int:
+            return 1
+    x = f({"x": 1})
+    reveal_type(x)  # revealed: int
+
+    # TODO: error: [invalid-argument-type] "Argument to function `f` is incorrect: Expected `T`, found `dict[str, int]`"
+    # we currently consider `TypedDict` instances to be subtypes of `dict`
+    f({"y": 1})
+```
+
+Diagnostics unrelated to the type-context are only reported once:
+
+```py
+def f[T](x: T) -> list[T]:
+    return [x]
+
+def a(x: list[bool], y: list[bool]): ...
+def b(x: list[int], y: list[int]): ...
+def c(x: list[int], y: list[int]): ...
+def _(x: int):
+    if x == 0:
+        y = a
+    elif x == 1:
+        y = b
+    else:
+        y = c
+
+    if x == 0:
+        z = True
+
+    y(f(True), [True])
+
+    # error: [possibly-unresolved-reference] "Name `z` used when possibly not defined"
+    y(f(True), [z])
+```

crates/ty_python_semantic/resources/mdtest/directives/assert_type.md

@@ -10,6 +10,7 @@ from typing_extensions import assert_type
 def _(x: int):
     assert_type(x, int)  # fine
     assert_type(x, str)  # error: [type-assertion-failure]
+    assert_type(assert_type(x, int), int)
 ```
 
 ## Narrowing

crates/ty_python_semantic/resources/mdtest/snapshots/assert_type.md_-_`assert_type`_-_Basic_(c507788da2659ec9).snap

@@ -17,6 +17,7 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/directives/assert_type.m
 3 | def _(x: int):
 4 |     assert_type(x, int)  # fine
 5 |     assert_type(x, str)  # error: [type-assertion-failure]
+6 |     assert_type(assert_type(x, int), int)
 ```
 
 # Diagnostics
@@ -31,6 +32,7 @@ error[type-assertion-failure]: Argument does not have asserted type `str`
   |     ^^^^^^^^^^^^-^^^^^^
   |                 |
   |                 Inferred type of argument is `int`
+6 |     assert_type(assert_type(x, int), int)
   |
 info: `str` and `int` are not equivalent types
 info: rule `type-assertion-failure` is enabled by default

crates/ty_python_semantic/resources/mdtest/typed_dict.md

@@ -152,7 +152,7 @@ Person(name="Alice")
 # error: [missing-typed-dict-key] "Missing required key 'age' in TypedDict `Person` constructor"
 Person({"name": "Alice"})
 
-# TODO: this should be an error, similar to the above
+# error: [missing-typed-dict-key] "Missing required key 'age' in TypedDict `Person` constructor"
 accepts_person({"name": "Alice"})
 # TODO: this should be an error, similar to the above
 house.owner = {"name": "Alice"}
@@ -171,7 +171,7 @@ Person(name=None, age=30)
 # error: [invalid-argument-type] "Invalid argument to key "name" with declared type `str` on TypedDict `Person`: value of type `None`"
 Person({"name": None, "age": 30})
 
-# TODO: this should be an error, similar to the above
+# error: [invalid-argument-type] "Invalid argument to key "name" with declared type `str` on TypedDict `Person`: value of type `None`"
 accepts_person({"name": None, "age": 30})
 # TODO: this should be an error, similar to the above
 house.owner = {"name": None, "age": 30}
@@ -190,7 +190,7 @@ Person(name="Alice", age=30, extra=True)
 # error: [invalid-key] "Invalid key access on TypedDict `Person`: Unknown key "extra""
 Person({"name": "Alice", "age": 30, "extra": True})
 
-# TODO: this should be an error
+# error: [invalid-key] "Invalid key access on TypedDict `Person`: Unknown key "extra""
 accepts_person({"name": "Alice", "age": 30, "extra": True})
 # TODO: this should be an error
 house.owner = {"name": "Alice", "age": 30, "extra": True}

crates/ty_python_semantic/src/types.rs

@@ -4194,20 +4194,26 @@ impl<'db> Type<'db> {
                     .into()
                 }
 
-                Some(KnownFunction::AssertType) => Binding::single(
-                    self,
-                    Signature::new(
-                        Parameters::new([
-                            Parameter::positional_only(Some(Name::new_static("value")))
-                                .with_annotated_type(Type::any()),
-                            Parameter::positional_only(Some(Name::new_static("type")))
-                                .type_form()
-                                .with_annotated_type(Type::any()),
-                        ]),
-                        Some(Type::none(db)),
-                    ),
-                )
-                .into(),
+                Some(KnownFunction::AssertType) => {
+                    let val_ty =
+                        BoundTypeVarInstance::synthetic(db, "T", TypeVarVariance::Invariant);
+
+                    Binding::single(
+                        self,
+                        Signature::new_generic(
+                            Some(GenericContext::from_typevar_instances(db, [val_ty])),
+                            Parameters::new([
+                                Parameter::positional_only(Some(Name::new_static("value")))
+                                    .with_annotated_type(Type::TypeVar(val_ty)),
+                                Parameter::positional_only(Some(Name::new_static("type")))
+                                    .type_form()
+                                    .with_annotated_type(Type::any()),
+                            ]),
+                            Some(Type::TypeVar(val_ty)),
+                        ),
+                    )
+                    .into()
+                }
 
                 Some(KnownFunction::AssertNever) => {
                     Binding::single(

crates/ty_python_semantic/src/types/builder.rs

@@ -1077,7 +1077,11 @@ impl<'db> InnerIntersectionBuilder<'db> {
             // don't need to worry about finding any particular constraint more than once.
             let constraints = constraints.elements(db);
             let mut positive_constraint_count = 0;
-            for positive in &self.positive {
+            for (i, positive) in self.positive.iter().enumerate() {
+                if i == typevar_index {
+                    continue;
+                }
+
                 // This linear search should be fine as long as we don't encounter typevars with
                 // thousands of constraints.
                 positive_constraint_count += constraints

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

@@ -33,10 +33,10 @@ use crate::types::{
     BoundMethodType, ClassLiteral, DataclassParams, FieldInstance, KnownBoundMethodType,
     KnownClass, KnownInstanceType, MemberLookupPolicy, PropertyInstanceType, SpecialFormType,
     TrackedConstraintSet, TypeAliasType, TypeContext, UnionBuilder, UnionType,
-    WrapperDescriptorKind, enums, ide_support, todo_type,
+    WrapperDescriptorKind, enums, ide_support, infer_isolated_expression, todo_type,
 };
 use ruff_db::diagnostic::{Annotation, Diagnostic, SubDiagnostic, SubDiagnosticSeverity};
-use ruff_python_ast::{self as ast, PythonVersion};
+use ruff_python_ast::{self as ast, ArgOrKeyword, PythonVersion};
 
 /// Binding information for a possible union of callables. At a call site, the arguments must be
 /// compatible with _all_ of the types in the union for the call to be valid.
@@ -1776,7 +1776,7 @@ impl<'db> CallableBinding<'db> {
     }
 
     /// Returns the index of the matching overload in the form of [`MatchingOverloadIndex`].
-    fn matching_overload_index(&self) -> MatchingOverloadIndex {
+    pub(crate) fn matching_overload_index(&self) -> MatchingOverloadIndex {
         let mut matching_overloads = self.matching_overloads();
         match matching_overloads.next() {
             None => MatchingOverloadIndex::None,
@@ -1794,8 +1794,15 @@ impl<'db> CallableBinding<'db> {
         }
     }
 
+    /// Returns all overloads for this call binding, including overloads that did not match.
+    pub(crate) fn overloads(&self) -> &[Binding<'db>] {
+        self.overloads.as_slice()
+    }
+
     /// Returns an iterator over all the overloads that matched for this call binding.
-    pub(crate) fn matching_overloads(&self) -> impl Iterator<Item = (usize, &Binding<'db>)> {
+    pub(crate) fn matching_overloads(
+        &self,
+    ) -> impl Iterator<Item = (usize, &Binding<'db>)> + Clone {
         self.overloads
             .iter()
             .enumerate()
@@ -2026,7 +2033,7 @@ enum OverloadCallReturnType<'db> {
 }
 
 #[derive(Debug)]
-enum MatchingOverloadIndex {
+pub(crate) enum MatchingOverloadIndex {
     /// No matching overloads found.
     None,
 
@@ -2504,9 +2511,17 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
         if let Some(return_ty) = self.signature.return_ty
             && let Some(call_expression_tcx) = self.call_expression_tcx.annotation
         {
-            // Ignore any specialization errors here, because the type context is only used to
-            // optionally widen the return type.
-            let _ = builder.infer(return_ty, call_expression_tcx);
+            match call_expression_tcx {
+                // A type variable is not a useful type-context for expression inference, and applying it
+                // to the return type can lead to confusing unions in nested generic calls.
+                Type::TypeVar(_) => {}
+
+                _ => {
+                    // Ignore any specialization errors here, because the type context is only used as a hint
+                    // to infer a more assignable return type.
+                    let _ = builder.infer(return_ty, call_expression_tcx);
+                }
+            }
         }
 
         let parameters = self.signature.parameters();
@@ -3289,6 +3304,23 @@ impl<'db> BindingError<'db> {
                     return;
                 };
 
+                // Re-infer the argument type of call expressions, ignoring the type context for more
+                // precise error messages.
+                let provided_ty = match Self::get_argument_node(node, *argument_index) {
+                    None => *provided_ty,
+
+                    // Ignore starred arguments, as those are difficult to re-infer.
+                    Some(
+                        ast::ArgOrKeyword::Arg(ast::Expr::Starred(_))
+                        | ast::ArgOrKeyword::Keyword(ast::Keyword { arg: None, .. }),
+                    ) => *provided_ty,
+
+                    Some(
+                        ast::ArgOrKeyword::Arg(value)
+                        | ast::ArgOrKeyword::Keyword(ast::Keyword { value, .. }),
+                    ) => infer_isolated_expression(context.db(), context.scope(), value),
+                };
+
                 let provided_ty_display = provided_ty.display(context.db());
                 let expected_ty_display = expected_ty.display(context.db());
 
@@ -3624,22 +3656,29 @@ impl<'db> BindingError<'db> {
         }
     }
 
-    fn get_node(node: ast::AnyNodeRef, argument_index: Option<usize>) -> ast::AnyNodeRef {
+    fn get_node(node: ast::AnyNodeRef<'_>, argument_index: Option<usize>) -> ast::AnyNodeRef<'_> {
         // If we have a Call node and an argument index, report the diagnostic on the correct
         // argument node; otherwise, report it on the entire provided node.
+        match Self::get_argument_node(node, argument_index) {
+            Some(ast::ArgOrKeyword::Arg(expr)) => expr.into(),
+            Some(ast::ArgOrKeyword::Keyword(expr)) => expr.into(),
+            None => node,
+        }
+    }
+
+    fn get_argument_node(
+        node: ast::AnyNodeRef<'_>,
+        argument_index: Option<usize>,
+    ) -> Option<ArgOrKeyword<'_>> {
         match (node, argument_index) {
-            (ast::AnyNodeRef::ExprCall(call_node), Some(argument_index)) => {
-                match call_node
+            (ast::AnyNodeRef::ExprCall(call_node), Some(argument_index)) => Some(
+                call_node
                     .arguments
                     .arguments_source_order()
                     .nth(argument_index)
-                    .expect("argument index should not be out of range")
-                {
-                    ast::ArgOrKeyword::Arg(expr) => expr.into(),
-                    ast::ArgOrKeyword::Keyword(keyword) => keyword.into(),
-                }
-            }
-            _ => node,
+                    .expect("argument index should not be out of range"),
+            ),
+            _ => None,
         }
     }
 }