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,23 @@ 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: 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: int | str = f(lst(b))
+    reveal_type(x4)  # revealed: int | str
+
+    x5: int | str | None = f(lst(b))
+    reveal_type(x5)  # revealed: int | str
 ```

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/class/super.md

@@ -206,8 +206,7 @@ class BuilderMeta2(type):
     ) -> BuilderMeta2:
         # revealed: <super: <class 'BuilderMeta2'>, <class 'BuilderMeta2'>>
         s = reveal_type(super())
-        # TODO: should be `BuilderMeta2` (needs https://github.com/astral-sh/ty/issues/501)
-        # revealed:  Unknown
+        # revealed:  BuilderMeta2
         return reveal_type(s.__new__(cls, name, bases, dct))
 
 class Foo[T]:

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__)