[red-knot] Allow explicit specialization of generic classes (#17023)

This PR lets you explicitly specialize a generic class using a subscript
expression. It introduces three new Rust types for representing classes:

- `NonGenericClass`
- `GenericClass` (not specialized)
- `GenericAlias` (specialized)

and two enum wrappers:

- `ClassType` (a non-generic class or generic alias, represents a class
_type_ at runtime)
- `ClassLiteralType` (a non-generic class or generic class, represents a
class body in the AST)

We also add internal support for specializing callables, in particular
function literals. (That is, the internal `Type` representation now
attaches an optional specialization to a function literal.) This is used
in this PR for the methods of a generic class, but should also give us
most of what we need for specializing generic _functions_ (which this PR
does not yet tackle).

---------

Co-authored-by: Alex Waygood <[email protected]>
Co-authored-by: Carl Meyer <[email protected]>

Author: 3 people

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

@@ -13,17 +13,13 @@ class C[T]: ...
 A class that inherits from a generic class, and fills its type parameters with typevars, is generic:
 
 ```py
-# TODO: no error
-# error: [non-subscriptable]
 class D[U](C[U]): ...
 ```
 
 A class that inherits from a generic class, but fills its type parameters with concrete types, is
 _not_ generic:
 
 ```py
-# TODO: no error
-# error: [non-subscriptable]
 class E(C[int]): ...
 ```
 
@@ -57,33 +53,85 @@ class D(C[T]): ...
 
 (Examples `E` and `F` from above do not have analogues in the legacy syntax.)
 
-## Inferring generic class parameters
+## Specializing generic classes explicitly
 
 The type parameter can be specified explicitly:
 
 ```py
 class C[T]:
     x: T
 
-# TODO: no error
-# TODO: revealed: C[int]
-# error: [non-subscriptable]
-reveal_type(C[int]())  # revealed: C
+reveal_type(C[int]())  # revealed: C[int]
+```
+
+The specialization must match the generic types:
+
+```py
+# error: [too-many-positional-arguments] "Too many positional arguments to class `C`: expected 1, got 2"
+reveal_type(C[int, int]())  # revealed: Unknown
+```
+
+If the type variable has an upper bound, the specialized type must satisfy that bound:
+
+```py
+class Bounded[T: int]: ...
+class BoundedByUnion[T: int | str]: ...
+class IntSubclass(int): ...
+
+reveal_type(Bounded[int]())  # revealed: Bounded[int]
+reveal_type(Bounded[IntSubclass]())  # revealed: Bounded[IntSubclass]
+
+# error: [invalid-argument-type] "Object of type `str` cannot be assigned to parameter 1 (`T`) of class `Bounded`; expected type `int`"
+reveal_type(Bounded[str]())  # revealed: Unknown
+
+# error:  [invalid-argument-type] "Object of type `int | str` cannot be assigned to parameter 1 (`T`) of class `Bounded`; expected type `int`"
+reveal_type(Bounded[int | str]())  # revealed: Unknown
+
+reveal_type(BoundedByUnion[int]())  # revealed: BoundedByUnion[int]
+reveal_type(BoundedByUnion[IntSubclass]())  # revealed: BoundedByUnion[IntSubclass]
+reveal_type(BoundedByUnion[str]())  # revealed: BoundedByUnion[str]
+reveal_type(BoundedByUnion[int | str]())  # revealed: BoundedByUnion[int | str]
 ```
 
+If the type variable is constrained, the specialized type must satisfy those constraints:
+
+```py
+class Constrained[T: (int, str)]: ...
+
+reveal_type(Constrained[int]())  # revealed: Constrained[int]
+
+# TODO: error: [invalid-argument-type]
+# TODO: revealed: Constrained[Unknown]
+reveal_type(Constrained[IntSubclass]())  # revealed: Constrained[IntSubclass]
+
+reveal_type(Constrained[str]())  # revealed: Constrained[str]
+
+# TODO: error: [invalid-argument-type]
+# TODO: revealed: Unknown
+reveal_type(Constrained[int | str]())  # revealed: Constrained[int | str]
+
+# error: [invalid-argument-type] "Object of type `object` cannot be assigned to parameter 1 (`T`) of class `Constrained`; expected type `int | str`"
+reveal_type(Constrained[object]())  # revealed: Unknown
+```
+
+## Inferring generic class parameters
+
 We can infer the type parameter from a type context:
 
 ```py
+class C[T]:
+    x: T
+
 c: C[int] = C()
 # TODO: revealed: C[int]
-reveal_type(c)  # revealed: C
+reveal_type(c)  # revealed: C[Unknown]
 ```
 
 The typevars of a fully specialized generic class should no longer be visible:
 
 ```py
 # TODO: revealed: int
-reveal_type(c.x)  # revealed: T
+reveal_type(c.x)  # revealed: Unknown
 ```
 
 If the type parameter is not specified explicitly, and there are no constraints that let us infer a
@@ -92,15 +140,13 @@ specific type, we infer the typevar's default type:
 ```py
 class D[T = int]: ...
 
-# TODO: revealed: D[int]
-reveal_type(D())  # revealed: D
+reveal_type(D())  # revealed: D[int]
 ```
 
 If a typevar does not provide a default, we use `Unknown`:
 
 ```py
-# TODO: revealed: C[Unknown]
-reveal_type(C())  # revealed: C
+reveal_type(C())  # revealed: C[Unknown]
 ```
 
 If the type of a constructor parameter is a class typevar, we can use that to infer the type
@@ -111,17 +157,14 @@ class E[T]:
     def __init__(self, x: T) -> None: ...
 
 # TODO: revealed: E[int] or E[Literal[1]]
-# TODO should not emit an error
-# error: [invalid-argument-type] "Object of type `Literal[1]` cannot be assigned to parameter 2 (`x`) of bound method `__init__`; expected type `T`"
-reveal_type(E(1))  # revealed: E
+reveal_type(E(1))  # revealed: E[Unknown]
 ```
 
 The types inferred from a type context and from a constructor parameter must be consistent with each
 other:
 
 ```py
-# TODO: the error should not leak the `T` typevar and should mention `E[int]`
-# error: [invalid-argument-type] "Object of type `Literal["five"]` cannot be assigned to parameter 2 (`x`) of bound method `__init__`; expected type `T`"
+# TODO: error: [invalid-argument-type]
 wrong_innards: E[int] = E("five")
 ```
 
@@ -134,17 +177,33 @@ propagate through:
 class Base[T]:
     x: T | None = None
 
-# TODO: no error
-# error: [non-subscriptable]
 class Sub[U](Base[U]): ...
 
+reveal_type(Base[int].x)  # revealed: int | None
+reveal_type(Sub[int].x)  # revealed: int | None
+```
+
+## Generic methods
+
+Generic classes can contain methods that are themselves generic. The generic methods can refer to
+the typevars of the enclosing generic class, and introduce new (distinct) typevars that are only in
+scope for the method.
+
+```py
+class C[T]:
+    def method[U](self, u: U) -> U:
+        return u
+    # error: [unresolved-reference]
+    def cannot_use_outside_of_method(self, u: U): ...
+
+    # TODO: error
+    def cannot_shadow_class_typevar[T](self, t: T): ...
+
+c: C[int] = C[int]()
 # TODO: no error
-# TODO: revealed: int | None
-# error: [non-subscriptable]
-reveal_type(Base[int].x)  # revealed: T | None
-# TODO: revealed: int | None
-# error: [non-subscriptable]
-reveal_type(Sub[int].x)  # revealed: T | None
+# TODO: revealed: str or Literal["string"]
+# error: [invalid-argument-type]
+reveal_type(c.method("string"))  # revealed: U
 ```
 
 ## Cyclic class definition
@@ -158,8 +217,6 @@ Here, `Sub` is not a generic class, since it fills its superclass's type paramet
 
 ```pyi
 class Base[T]: ...
-# TODO: no error
-# error: [non-subscriptable]
 class Sub(Base[Sub]): ...
 
 reveal_type(Sub)  # revealed: Literal[Sub]
@@ -171,9 +228,6 @@ A similar case can work in a non-stub file, if forward references are stringifie
 
 ```py
 class Base[T]: ...
-
-# TODO: no error
-# error: [non-subscriptable]
 class Sub(Base["Sub"]): ...
 
 reveal_type(Sub)  # revealed: Literal[Sub]
@@ -186,8 +240,6 @@ In a non-stub file, without stringified forward references, this raises a `NameE
 ```py
 class Base[T]: ...
 
-# TODO: the unresolved-reference error is correct, the non-subscriptable is not
-# error: [non-subscriptable]
 # error: [unresolved-reference]
 class Sub(Base[Sub]): ...
 ```

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

@@ -82,18 +82,51 @@ class C[T]:
     def m2(self, x: T) -> T:
         return x
 
-c: C[int] = C()
-# TODO: no error
-# error: [invalid-argument-type]
+c: C[int] = C[int]()
 c.m1(1)
-# TODO: no error
-# error: [invalid-argument-type]
 c.m2(1)
-# TODO: expected type `int`
-# error: [invalid-argument-type] "Object of type `Literal["string"]` cannot be assigned to parameter 2 (`x`) of bound method `m2`; expected type `T`"
+# error: [invalid-argument-type] "Object of type `Literal["string"]` cannot be assigned to parameter 2 (`x`) of bound method `m2`; expected type `int`"
 c.m2("string")
 ```
 
+## Functions on generic classes are descriptors
+
+This repeats the tests in the [Functions as descriptors](./call/methods.md) test suite, but on a
+generic class. This ensures that we are carrying any specializations through the entirety of the
+descriptor protocol, which is how `self` parameters are bound to instance methods.
+
+```py
+from inspect import getattr_static
+
+class C[T]:
+    def f(self, x: T) -> str:
+        return "a"
+
+reveal_type(getattr_static(C[int], "f"))  # revealed: Literal[f[int]]
+reveal_type(getattr_static(C[int], "f").__get__)  # revealed: <method-wrapper `__get__` of `f[int]`>
+reveal_type(getattr_static(C[int], "f").__get__(None, C[int]))  # revealed: Literal[f[int]]
+# revealed: <bound method `f` of `C[int]`>
+reveal_type(getattr_static(C[int], "f").__get__(C[int](), C[int]))
+
+reveal_type(C[int].f)  # revealed: Literal[f[int]]
+reveal_type(C[int]().f)  # revealed: <bound method `f` of `C[int]`>
+
+bound_method = C[int]().f
+reveal_type(bound_method.__self__)  # revealed: C[int]
+reveal_type(bound_method.__func__)  # revealed: Literal[f[int]]
+
+reveal_type(C[int]().f(1))  # revealed: str
+reveal_type(bound_method(1))  # revealed: str
+
+C[int].f(1)  # error: [missing-argument]
+reveal_type(C[int].f(C[int](), 1))  # revealed: str
+
+class D[U](C[U]):
+    pass
+
+reveal_type(D[int]().f)  # revealed: <bound method `f` of `D[int]`>
+```
+
 ## Methods can mention other typevars
 
 > A type variable used in a method that does not match any of the variables that parameterize the
@@ -127,7 +160,6 @@ c: C[int] = C()
 # TODO: no errors
 # TODO: revealed: str
 # error: [invalid-argument-type]
-# error: [invalid-argument-type]
 reveal_type(c.m(1, "string"))  # revealed: S
 ```
 

crates/red_knot_python_semantic/resources/mdtest/narrow/type.md

@@ -109,6 +109,24 @@ def _(x: A | B):
         reveal_type(x)  # revealed: A
 ```
 
+## Narrowing for generic classes
+
+Note that `type` returns the runtime class of an object, which does _not_ include specializations in
+the case of a generic class. (The typevars are erased.) That means we cannot narrow the type to the
+specialization that we compare with; we must narrow to an unknown specialization of the generic
+class.
+
+```py
+class A[T = int]: ...
+class B: ...
+
+def _[T](x: A | B):
+    if type(x) is A[str]:
+        reveal_type(x)  # revealed: A[int] & A[Unknown] | B & A[Unknown]
+    else:
+        reveal_type(x)  # revealed: A[int] | B
+```
+
 ## Limitations
 
 ```py

crates/red_knot_python_semantic/resources/mdtest/stubs/class.md

@@ -8,13 +8,10 @@ In type stubs, classes can reference themselves in their base class definitions.
 ```pyi
 class Foo[T]: ...
 
-# TODO: actually is subscriptable
-# error: [non-subscriptable]
 class Bar(Foo[Bar]): ...
 
 reveal_type(Bar)  # revealed: Literal[Bar]
-# TODO: Instead of `Literal[Foo]`, we might eventually want to show a type that involves the type parameter.
-reveal_type(Bar.__mro__)  # revealed: tuple[Literal[Bar], Literal[Foo], Literal[object]]
+reveal_type(Bar.__mro__)  # revealed: tuple[Literal[Bar], Literal[Foo[Bar]], Literal[object]]
 ```
 
 ## Access to attributes declared in stubs

crates/red_knot_python_semantic/src/symbol.rs

@@ -425,6 +425,17 @@ impl<'db> SymbolAndQualifiers<'db> {
         self.qualifiers.contains(TypeQualifiers::CLASS_VAR)
     }
 
+    #[must_use]
+    pub(crate) fn map_type(
+        self,
+        f: impl FnOnce(Type<'db>) -> Type<'db>,
+    ) -> SymbolAndQualifiers<'db> {
+        SymbolAndQualifiers {
+            symbol: self.symbol.map_type(f),
+            qualifiers: self.qualifiers,
+        }
+    }
+
     /// Transform symbol and qualifiers into a [`LookupResult`],
     /// a [`Result`] type in which the `Ok` variant represents a definitely bound symbol
     /// and the `Err` variant represents a symbol that is either definitely or possibly unbound.