improve generic call expression inference

Author: ibraheemdev

crates/ty_python_semantic/resources/mdtest/assignment/annotations.md

@@ -427,14 +427,13 @@ a = f("a")
 reveal_type(a)  # revealed: list[Literal["a"]]
 
 b: list[int | Literal["a"]] = f("a")
-reveal_type(b)  # revealed: list[Literal["a"] | int]
+reveal_type(b)  # revealed: list[int | Literal["a"]]
 
 c: list[int | str] = f("a")
-reveal_type(c)  # revealed: list[str | int]
+reveal_type(c)  # revealed: list[int | str]
 
 d: list[int | tuple[int, int]] = f((1, 2))
-# TODO: We could avoid reordering the union elements here.
-reveal_type(d)  # revealed: list[tuple[int, int] | int]
+reveal_type(d)  # revealed: list[int | tuple[int, int]]
 
 e: list[int] = f(True)
 reveal_type(e)  # revealed: list[int]
@@ -455,7 +454,120 @@ j: int | str = f2(True)
 reveal_type(j)  # revealed: Literal[True]
 ```
 
-Types are not widened unnecessarily:
+The function arguments are also inferred using the type context:
+
+```py
+from typing import TypedDict
+
+class TD(TypedDict):
+    x: int
+
+def f[T](x: list[T]) -> T:
+    return x[0]
+
+a: TD = f([{"x": 0}, {"x": 1}])
+reveal_type(a)  # revealed: TD
+
+b: TD | None = f([{"x": 0}, {"x": 1}])
+reveal_type(b)  # revealed: TD
+
+# error: [missing-typed-dict-key] "Missing required key 'x' in TypedDict `TD` constructor"
+# error: [invalid-key] "Invalid key for TypedDict `TD`: Unknown key "y""
+# error: [invalid-assignment] "Object of type `Unknown | dict[Unknown | str, Unknown | int]` is not assignable to `TD`"
+c: TD = f([{"y": 0}, {"x": 1}])
+
+# error: [missing-typed-dict-key] "Missing required key 'x' in TypedDict `TD` constructor"
+# error: [invalid-key] "Invalid key for TypedDict `TD`: Unknown key "y""
+# error: [invalid-assignment] "Object of type `Unknown | dict[Unknown | str, Unknown | int]` is not assignable to `TD | None`"
+c: TD | None = f([{"y": 0}, {"x": 1}])
+```
+
+## Prefer the declared type of generic classes
+
+```toml
+[environment]
+python-version = "3.12"
+```
+
+```py
+from typing import Any
+
+def f[T](x: T) -> list[T]:
+    return [x]
+
+def f2[T](x: T) -> list[T] | None:
+    return [x]
+
+def f3[T](x: T) -> list[T] | dict[T, T]:
+    return [x]
+
+a = f(1)
+reveal_type(a)  # revealed: list[Literal[1]]
+
+b: list[Any] = f(1)
+reveal_type(b)  # revealed: list[Any]
+
+c: list[Any] = [1]
+reveal_type(c)  # revealed: list[Any]
+
+d: list[Any] | None = f(1)
+reveal_type(d)  # revealed: list[Any]
+
+e: list[Any] | None = [1]
+reveal_type(e)  # revealed: list[Any]
+
+f: list[Any] | None = f2(1)
+# TODO: Better constraint solver.
+reveal_type(f)  # revealed: list[Literal[1]] | None
+
+g: list[Any] | dict[Any, Any] = f3(1)
+# TODO: Better constraint solver.
+reveal_type(g)  # revealed: list[Literal[1]] | dict[Literal[1], Literal[1]]
+```
+
+## Narrow generic unions
+
+```toml
+[environment]
+python-version = "3.12"
+```
+
+```py
+from typing import reveal_type, TypedDict
+
+def id[T](x: T) -> T:
+    return x
+
+def _(narrow: dict[str, str], target: list[str] | dict[str, str] | None):
+    target = id(narrow)
+    reveal_type(target)  # revealed: dict[str, str]
+
+def _(narrow: list[str], target: list[str] | dict[str, str] | None):
+    target = id(narrow)
+    reveal_type(target)  # revealed: list[str]
+
+def _(narrow: list[str] | dict[str, str], target: list[str] | dict[str, str] | None):
+    target = id(narrow)
+    reveal_type(target)  # revealed: list[str] | dict[str, str]
+
+class TD(TypedDict):
+    x: int
+
+def _(target: list[TD] | dict[str, TD] | None):
+    target = id([{"x": 1}])
+    reveal_type(target)  # revealed: list[TD]
+
+def _(target: list[TD] | dict[str, TD] | None):
+    target = id({"x": {"x": 1}})
+    reveal_type(target)  # revealed: dict[str, TD]
+```
+
+## Prefer the inferred type of non-generic classes
+
+```toml
+[environment]
+python-version = "3.12"
+```
 
 ```py
 def id[T](x: T) -> T:
@@ -476,10 +588,8 @@ def _(i: int):
     b: list[int | None] | None = id([i])
     c: list[int | None] | int | None = id([i])
     reveal_type(a)  # revealed: list[int | None]
-    # TODO: these should reveal `list[int | None]`
-    # we currently do not use the call expression annotation as type context for argument inference
-    reveal_type(b)  # revealed: list[Unknown | int]
-    reveal_type(c)  # revealed: list[Unknown | int]
+    reveal_type(b)  # revealed: list[int | None]
+    reveal_type(c)  # revealed: list[int | None]
 
     a: list[int | None] | None = [i]
     b: list[int | None] | None = lst(i)
@@ -494,4 +604,39 @@ def _(i: int):
     reveal_type(a)  # revealed: list[Unknown]
     reveal_type(b)  # revealed: list[Unknown]
     reveal_type(c)  # revealed: list[Unknown]
+
+def f[T](x: list[T]) -> T:
+    return x[0]
+
+def _(a: int, b: str, c: int | str):
+    x1: int = f(lst(a))
+    reveal_type(x1)  # revealed: int
+
+    x2: int | str = f(lst(a))
+    reveal_type(x2)  # revealed: int
+
+    x3: int | None = f(lst(a))
+    reveal_type(x3)  # revealed: int
+
+    x4: str = f(lst(b))
+    reveal_type(x4)  # revealed: str
+
+    x5: int | str = f(lst(b))
+    reveal_type(x5)  # revealed: str
+
+    x6: str | None = f(lst(b))
+    reveal_type(x6)  # revealed: str
+
+    x7: int | str = f(lst(c))
+    reveal_type(x7)  # revealed: int | str
+
+    x8: int | str = f(lst(c))
+    reveal_type(x8)  # revealed: int | str
+
+    # TODO: Ideally this would reveal `int | str`. This is a known limitation of our
+    # call inference solver, and would # require an extra inference attempt without type
+    # context, or with type context # of subsets of the union, both of which are impractical
+    # for performance reasons.
+    x9: int | str | None = f(lst(c))
+    reveal_type(x9)  # revealed: int | str | None
 ```

crates/ty_python_semantic/resources/mdtest/bidirectional.md

@@ -50,8 +50,8 @@ def _(l: list[int] | None = None):
 def f[T](x: T, cond: bool) -> T | list[T]:
     return x if cond else [x]
 
-# TODO: no error
-# error: [invalid-assignment] "Object of type `Literal[1] | list[Literal[1]]` is not assignable to `int | list[int]`"
+# TODO: Better constraint solver.
+# error: [invalid-assignment]
 l5: int | list[int] = f(1, True)
 ```
 

crates/ty_python_semantic/resources/mdtest/dataclasses/fields.md

@@ -37,7 +37,7 @@ class Data:
     content: list[int] = field(default_factory=list)
     timestamp: datetime = field(default_factory=datetime.now, init=False)
 
-# revealed: (self: Data, content: list[int] = Unknown) -> None
+# revealed: (self: Data, content: list[int] = list[int]) -> None
 reveal_type(Data.__init__)
 
 data = Data([1, 2, 3])
@@ -63,7 +63,6 @@ class Person:
     age: int | None = field(default=None, kw_only=True)
     role: str = field(default="user", kw_only=True)
 
-# TODO: this would ideally show a default value of `None` for `age`
 # revealed: (self: Person, name: str, *, age: int | None = None, role: str = Literal["user"]) -> None
 reveal_type(Person.__init__)
 

crates/ty_python_semantic/src/types.rs

@@ -883,21 +883,31 @@ impl<'db> Type<'db> {
             _ => false,
         }
     }
-
     // If the type is a specialized instance of the given `KnownClass`, returns the specialization.
     pub(crate) fn known_specialization(
         &self,
         db: &'db dyn Db,
         known_class: KnownClass,
     ) -> Option<Specialization<'db>> {
         let class_literal = known_class.try_to_class_literal(db)?;
-        self.specialization_of(db, Some(class_literal))
+        self.specialization_of(db, class_literal)
+    }
+
+    // If this type is a class instance, returns its specialization.
+    pub(crate) fn class_specialization(self, db: &'db dyn Db) -> Option<Specialization<'db>> {
+        self.specialization_of_optional(db, None)
     }
 
     // If the type is a specialized instance of the given class, returns the specialization.
-    //
-    // If no class is provided, returns the specialization of any class instance.
     pub(crate) fn specialization_of(
+        self,
+        db: &'db dyn Db,
+        expected_class: ClassLiteral<'_>,
+    ) -> Option<Specialization<'db>> {
+        self.specialization_of_optional(db, Some(expected_class))
+    }
+
+    fn specialization_of_optional(
         self,
         db: &'db dyn Db,
         expected_class: Option<ClassLiteral<'_>>,
@@ -5551,7 +5561,7 @@ impl<'db> Type<'db> {
     ) -> Result<Bindings<'db>, CallError<'db>> {
         self.bindings(db)
             .match_parameters(db, argument_types)
-            .check_types(db, argument_types, &TypeContext::default(), &[])
+            .check_types(db, argument_types, TypeContext::default(), &[])
     }
 
     /// Look up a dunder method on the meta-type of `self` and call it.
@@ -5603,7 +5613,8 @@ impl<'db> Type<'db> {
                 let bindings = dunder_callable
                     .bindings(db)
                     .match_parameters(db, argument_types)
-                    .check_types(db, argument_types, &tcx, &[])?;
+                    .check_types(db, argument_types, tcx, &[])?;
+
                 if boundness == Definedness::PossiblyUndefined {
                     return Err(CallDunderError::PossiblyUnbound(Box::new(bindings)));
                 }