[red-knot] Super-basic generic inference at call sites (#17301)

This PR adds **_very_** basic inference of generic typevars at call
sites. It does not bring in a full unification algorithm, and there are
a few TODOs in the test suite that are not discharged by this. But it
handles a good number of useful cases! And the PR does not add anything
that would go away with a more sophisticated constraint solver.

In short, we just look for typevars in the formal parameters, and assume
that the inferred type of the corresponding argument is what that
typevar should map to. If a typevar appears more than once, we union
together the corresponding argument types.

Cases we are not yet handling:

- We are not widening literals.
- We are not recursing into parameters that are themselves generic
aliases.
- We are not being very clever with parameters that are union types.

---------

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>
Co-authored-by: Carl Meyer <carl@astral.sh>

Author: 3 people

crates/red_knot_python_semantic/resources/mdtest/call/methods.md

@@ -410,29 +410,19 @@ def does_nothing[T](f: T) -> T:
 
 class C:
     @classmethod
-    # TODO: no error should be emitted here (needs support for generics)
-    # error: [invalid-argument-type]
     @does_nothing
     def f1(cls: type[C], x: int) -> str:
         return "a"
-    # TODO: no error should be emitted here (needs support for generics)
-    # error: [invalid-argument-type]
+
     @does_nothing
     @classmethod
     def f2(cls: type[C], x: int) -> str:
         return "a"
 
-# TODO: All of these should be `str` (and not emit an error), once we support generics
-
-# error: [call-non-callable]
-reveal_type(C.f1(1))  # revealed: Unknown
-# error: [call-non-callable]
-reveal_type(C().f1(1))  # revealed: Unknown
-
-# error: [call-non-callable]
-reveal_type(C.f2(1))  # revealed: Unknown
-# error: [call-non-callable]
-reveal_type(C().f2(1))  # revealed: Unknown
+reveal_type(C.f1(1))  # revealed: str
+reveal_type(C().f1(1))  # revealed: str
+reveal_type(C.f2(1))  # revealed: str
+reveal_type(C().f2(1))  # revealed: str
 ```
 
 [functions and methods]: https://docs.python.org/3/howto/descriptor.html#functions-and-methods

crates/red_knot_python_semantic/resources/mdtest/generics/classes.md

@@ -149,23 +149,102 @@ If a typevar does not provide a default, we use `Unknown`:
 reveal_type(C())  # revealed: C[Unknown]
 ```
 
+## Inferring generic class parameters from constructors
+
 If the type of a constructor parameter is a class typevar, we can use that to infer the type
-parameter:
+parameter. The types inferred from a type context and from a constructor parameter must be
+consistent with each other.
+
+## `__new__` only
 
 ```py
-class E[T]:
+class C[T]:
+    def __new__(cls, x: T) -> "C"[T]:
+        return object.__new__(cls)
+
+reveal_type(C(1))  # revealed: C[Literal[1]]
+
+# TODO: error: [invalid-argument-type]
+wrong_innards: C[int] = C("five")
+```
+
+## `__init__` only
+
+```py
+class C[T]:
     def __init__(self, x: T) -> None: ...
 
-# TODO: revealed: E[int] or E[Literal[1]]
-reveal_type(E(1))  # revealed: E[Unknown]
+reveal_type(C(1))  # revealed: C[Literal[1]]
+
+# TODO: error: [invalid-argument-type]
+wrong_innards: C[int] = C("five")
 ```
 
-The types inferred from a type context and from a constructor parameter must be consistent with each
-other:
+## Identical `__new__` and `__init__` signatures
 
 ```py
+class C[T]:
+    def __new__(cls, x: T) -> "C"[T]:
+        return object.__new__(cls)
+
+    def __init__(self, x: T) -> None: ...
+
+reveal_type(C(1))  # revealed: C[Literal[1]]
+
 # TODO: error: [invalid-argument-type]
-wrong_innards: E[int] = E("five")
+wrong_innards: C[int] = C("five")
+```
+
+## Compatible `__new__` and `__init__` signatures
+
+```py
+class C[T]:
+    def __new__(cls, *args, **kwargs) -> "C"[T]:
+        return object.__new__(cls)
+
+    def __init__(self, x: T) -> None: ...
+
+reveal_type(C(1))  # revealed: C[Literal[1]]
+
+# TODO: error: [invalid-argument-type]
+wrong_innards: C[int] = C("five")
+
+class D[T]:
+    def __new__(cls, x: T) -> "D"[T]:
+        return object.__new__(cls)
+
+    def __init__(self, *args, **kwargs) -> None: ...
+
+reveal_type(D(1))  # revealed: D[Literal[1]]
+
+# TODO: error: [invalid-argument-type]
+wrong_innards: D[int] = D("five")
+```
+
+## `__init__` is itself generic
+
+TODO: These do not currently work yet, because we don't correctly model the nested generic contexts.
+
+```py
+class C[T]:
+    def __init__[S](self, x: T, y: S) -> None: ...
+
+# TODO: no error
+# TODO: revealed: C[Literal[1]]
+# error: [invalid-argument-type]
+reveal_type(C(1, 1))  # revealed: C[Unknown]
+# TODO: no error
+# TODO: revealed: C[Literal[1]]
+# error: [invalid-argument-type]
+reveal_type(C(1, "string"))  # revealed: C[Unknown]
+# TODO: no error
+# TODO: revealed: C[Literal[1]]
+# error: [invalid-argument-type]
+reveal_type(C(1, True))  # revealed: C[Unknown]
+
+# TODO: error for the correct reason
+# error: [invalid-argument-type] "Argument to this function is incorrect: Expected `S`, found `Literal[1]`"
+wrong_innards: C[int] = C("five", 1)
 ```
 
 ## Generic subclass
@@ -200,10 +279,7 @@ class C[T]:
     def cannot_shadow_class_typevar[T](self, t: T): ...
 
 c: C[int] = C[int]()
-# TODO: no error
-# TODO: revealed: str or Literal["string"]
-# error: [invalid-argument-type]
-reveal_type(c.method("string"))  # revealed: U
+reveal_type(c.method("string"))  # revealed: Literal["string"]
 ```
 
 ## Cyclic class definition

crates/red_knot_python_semantic/resources/mdtest/generics/functions.md

@@ -43,33 +43,14 @@ def absurd[T]() -> T:
 If the type of a generic function parameter is a typevar, then we can infer what type that typevar
 is bound to at each call site.
 
-TODO: Note that some of the TODO revealed types have two options, since we haven't decided yet
-whether we want to infer a more specific `Literal` type where possible, or use heuristics to weaken
-the inferred type to e.g. `int`.
-
 ```py
 def f[T](x: T) -> T:
     return x
 
-# TODO: no error
-# TODO: revealed: int or Literal[1]
-# error: [invalid-argument-type]
-reveal_type(f(1))  # revealed: T
-
-# TODO: no error
-# TODO: revealed: float
-# error: [invalid-argument-type]
-reveal_type(f(1.0))  # revealed: T
-
-# TODO: no error
-# TODO: revealed: bool or Literal[true]
-# error: [invalid-argument-type]
-reveal_type(f(True))  # revealed: T
-
-# TODO: no error
-# TODO: revealed: str or Literal["string"]
-# error: [invalid-argument-type]
-reveal_type(f("string"))  # revealed: T
+reveal_type(f(1))  # revealed: Literal[1]
+reveal_type(f(1.0))  # revealed: float
+reveal_type(f(True))  # revealed: Literal[True]
+reveal_type(f("string"))  # revealed: Literal["string"]
 ```
 
 ## Inferring “deep” generic parameter types
@@ -82,7 +63,7 @@ def f[T](x: list[T]) -> T:
     return x[0]
 
 # TODO: revealed: float
-reveal_type(f([1.0, 2.0]))  # revealed: T
+reveal_type(f([1.0, 2.0]))  # revealed: Unknown
 ```
 
 ## Typevar constraints
@@ -93,7 +74,6 @@ in the function.
 
 ```py
 def good_param[T: int](x: T) -> None:
-    # TODO: revealed: T & int
     reveal_type(x)  # revealed: T
 ```
 
@@ -162,61 +142,41 @@ parameters simultaneously.
 def two_params[T](x: T, y: T) -> T:
     return x
 
-# TODO: no error
-# TODO: revealed: str
-# error: [invalid-argument-type]
-# error: [invalid-argument-type]
-reveal_type(two_params("a", "b"))  # revealed: T
-
-# TODO: no error
-# TODO: revealed: str | int
-# error: [invalid-argument-type]
-# error: [invalid-argument-type]
-reveal_type(two_params("a", 1))  # revealed: T
+reveal_type(two_params("a", "b"))  # revealed: Literal["a", "b"]
+reveal_type(two_params("a", 1))  # revealed: Literal["a", 1]
 ```
 
-```py
-def param_with_union[T](x: T | int, y: T) -> T:
-    return y
+When one of the parameters is a union, we attempt to find the smallest specialization that satisfies
+all of the constraints.
 
-# TODO: no error
-# TODO: revealed: str
-# error: [invalid-argument-type]
-reveal_type(param_with_union(1, "a"))  # revealed: T
+```py
+def union_param[T](x: T | None) -> T:
+    if x is None:
+        raise ValueError
+    return x
 
-# TODO: no error
-# TODO: revealed: str
-# error: [invalid-argument-type]
-# error: [invalid-argument-type]
-reveal_type(param_with_union("a", "a"))  # revealed: T
+reveal_type(union_param("a"))  # revealed: Literal["a"]
+reveal_type(union_param(1))  # revealed: Literal[1]
+reveal_type(union_param(None))  # revealed: Unknown
+```
 
-# TODO: no error
-# TODO: revealed: int
-# error: [invalid-argument-type]
-reveal_type(param_with_union(1, 1))  # revealed: T
+```py
+def union_and_nonunion_params[T](x: T | int, y: T) -> T:
+    return y
 
-# TODO: no error
-# TODO: revealed: str | int
-# error: [invalid-argument-type]
-# error: [invalid-argument-type]
-reveal_type(param_with_union("a", 1))  # revealed: T
+reveal_type(union_and_nonunion_params(1, "a"))  # revealed: Literal["a"]
+reveal_type(union_and_nonunion_params("a", "a"))  # revealed: Literal["a"]
+reveal_type(union_and_nonunion_params(1, 1))  # revealed: Literal[1]
+reveal_type(union_and_nonunion_params(3, 1))  # revealed: Literal[1]
+reveal_type(union_and_nonunion_params("a", 1))  # revealed: Literal["a", 1]
 ```
 
 ```py
 def tuple_param[T, S](x: T | S, y: tuple[T, S]) -> tuple[T, S]:
     return y
 
-# TODO: no error
-# TODO: revealed: tuple[str, int]
-# error: [invalid-argument-type]
-# error: [invalid-argument-type]
-reveal_type(tuple_param("a", ("a", 1)))  # revealed: tuple[T, S]
-
-# TODO: no error
-# TODO: revealed: tuple[str, int]
-# error: [invalid-argument-type]
-# error: [invalid-argument-type]
-reveal_type(tuple_param(1, ("a", 1)))  # revealed: tuple[T, S]
+reveal_type(tuple_param("a", ("a", 1)))  # revealed: tuple[Literal["a"], Literal[1]]
+reveal_type(tuple_param(1, ("a", 1)))  # revealed: tuple[Literal["a"], Literal[1]]
 ```
 
 ## Inferring nested generic function calls
@@ -231,15 +191,6 @@ def f[T](x: T) -> tuple[T, int]:
 def g[T](x: T) -> T | None:
     return x
 
-# TODO: no error
-# TODO: revealed: tuple[str | None, int]
-# error: [invalid-argument-type]
-# error: [invalid-argument-type]
-reveal_type(f(g("a")))  # revealed: tuple[T, int]
-
-# TODO: no error
-# TODO: revealed: tuple[str, int] | None
-# error: [invalid-argument-type]
-# error: [invalid-argument-type]
-reveal_type(g(f("a")))  # revealed: T | None
+reveal_type(f(g("a")))  # revealed: tuple[Literal["a"] | None, int]
+reveal_type(g(f("a")))  # revealed: tuple[Literal["a"], int] | None
 ```

crates/red_knot_python_semantic/resources/mdtest/generics/scoping.md

@@ -59,14 +59,8 @@ to a different type each time.
 def f[T](x: T) -> T:
     return x
 
-# TODO: no error
-# TODO: revealed: int or Literal[1]
-# error: [invalid-argument-type]
-reveal_type(f(1))  # revealed: T
-# TODO: no error
-# TODO: revealed: str or Literal["a"]
-# error: [invalid-argument-type]
-reveal_type(f("a"))  # revealed: T
+reveal_type(f(1))  # revealed: Literal[1]
+reveal_type(f("a"))  # revealed: Literal["a"]
 ```
 
 ## Methods can mention class typevars
@@ -157,10 +151,7 @@ class C[T]:
         return y
 
 c: C[int] = C()
-# TODO: no errors
-# TODO: revealed: str
-# error: [invalid-argument-type]
-reveal_type(c.m(1, "string"))  # revealed: S
+reveal_type(c.m(1, "string"))  # revealed: Literal["string"]
 ```
 
 ## Unbound typevars