[ty] Improve generic call expression inference

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

@@ -417,6 +417,8 @@ reveal_type(x) # revealed: Literal[1]
 python-version = "3.12"
 ```
 
+`generic_list.py`:
+
 ```py
 from typing import Literal
 
@@ -427,14 +429,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,31 +456,237 @@ j: int | str = f2(True)
 reveal_type(j)  # revealed: Literal[True]
 ```
 
-Types are not widened unnecessarily:
+A function's arguments are also inferred using the type context:
+
+`typed_dict.py`:
+
+```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}])
+```
+
+But not in a way that leads to assignability errors:
+
+`dict_any.py`:
 
 ```py
-def id[T](x: T) -> T:
+from typing import TypedDict, Any
+
+class TD(TypedDict, total=False):
+    x: str
+
+class TD2(TypedDict):
+    x: str
+
+def f(self, dt: dict[str, Any], key: str):
+    # TODO: This should not error once typed dict assignability is implemented.
+    # error: [invalid-assignment]
+    x1: TD = dt.get(key, {})
+    reveal_type(x1)  # revealed: TD
+
+    x2: TD = dt.get(key, {"x": 0})
+    reveal_type(x2)  # revealed: Any
+
+    x3: TD | None = dt.get(key, {})
+    # TODO: This should reveal `Any` once typed dict assignability is implemented.
+    reveal_type(x3)  # revealed: Any | None
+
+    x4: TD | None = dt.get(key, {"x": 0})
+    reveal_type(x4)  # revealed: Any
+
+    x5: TD2 = dt.get(key, {})
+    reveal_type(x5)  # revealed: Any
+
+    x6: TD2 = dt.get(key, {"x": 0})
+    reveal_type(x6)  # revealed: Any
+
+    x7: TD2 | None = dt.get(key, {})
+    reveal_type(x7)  # revealed: Any
+
+    x8: TD2 | None = dt.get(key, {"x": 0})
+    reveal_type(x8)  # revealed: Any
+```
+
+## 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]]
+```
+
+We currently prefer the generic declared type regardless of its variance:
+
+```py
+class Bivariant[T]:
+    pass
+
+class Covariant[T]:
+    def pop(self) -> T:
+        raise NotImplementedError
+
+class Contravariant[T]:
+    def push(self, value: T) -> None:
+        pass
+
+class Invariant[T]:
+    x: T
+
+def bivariant[T](x: T) -> Bivariant[T]:
+    return Bivariant()
+
+def covariant[T](x: T) -> Covariant[T]:
+    return Covariant()
+
+def contravariant[T](x: T) -> Contravariant[T]:
+    return Contravariant()
+
+def invariant[T](x: T) -> Invariant[T]:
+    return Invariant()
+
+x1 = bivariant(1)
+x2 = covariant(1)
+x3 = contravariant(1)
+x4 = invariant(1)
+
+reveal_type(x1)  # revealed: Bivariant[Literal[1]]
+reveal_type(x2)  # revealed: Covariant[Literal[1]]
+reveal_type(x3)  # revealed: Contravariant[Literal[1]]
+reveal_type(x4)  # revealed: Invariant[Literal[1]]
+
+x5: Bivariant[Any] = bivariant(1)
+x6: Covariant[Any] = covariant(1)
+x7: Contravariant[Any] = contravariant(1)
+x8: Invariant[Any] = invariant(1)
+
+# TODO: This could reveal `Bivariant[Any]`.
+reveal_type(x5)  # revealed: Bivariant[Literal[1]]
+reveal_type(x6)  # revealed: Covariant[Any]
+reveal_type(x7)  # revealed: Contravariant[Any]
+reveal_type(x8)  # revealed: Invariant[Any]
+```
+
+## Narrow generic unions
+
+```toml
+[environment]
+python-version = "3.12"
+```
+
+```py
+from typing import reveal_type, TypedDict
+
+def identity[T](x: T) -> T:
+    return x
+
+def _(narrow: dict[str, str], target: list[str] | dict[str, str] | None):
+    target = identity(narrow)
+    reveal_type(target)  # revealed: dict[str, str]
+
+def _(narrow: list[str], target: list[str] | dict[str, str] | None):
+    target = identity(narrow)
+    reveal_type(target)  # revealed: list[str]
+
+def _(narrow: list[str] | dict[str, str], target: list[str] | dict[str, str] | None):
+    target = identity(narrow)
+    reveal_type(target)  # revealed: list[str] | dict[str, str]
+
+class TD(TypedDict):
+    x: int
+
+def _(target: list[TD] | dict[str, TD] | None):
+    target = identity([{"x": 1}])
+    reveal_type(target)  # revealed: list[TD]
+
+def _(target: list[TD] | dict[str, TD] | None):
+    target = identity({"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 identity[T](x: T) -> T:
     return x
 
 def lst[T](x: T) -> list[T]:
     return [x]
 
 def _(i: int):
     a: int | None = i
-    b: int | None = id(i)
-    c: int | str | None = id(i)
+    b: int | None = identity(i)
+    c: int | str | None = identity(i)
     reveal_type(a)  # revealed: int
     reveal_type(b)  # revealed: int
     reveal_type(c)  # revealed: int
 
     a: list[int | None] | None = [i]
-    b: list[int | None] | None = id([i])
-    c: list[int | None] | int | None = id([i])
+    b: list[int | None] | None = identity([i])
+    c: list[int | None] | int | None = identity([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)
@@ -489,9 +696,44 @@ def _(i: int):
     reveal_type(c)  # revealed: list[int | None]
 
     a: list | None = []
-    b: list | None = id([])
-    c: list | int | None = id([])
+    b: list | None = identity([])
+    c: list | int | None = identity([])
     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__)