[red-knot] Use the right scope when considering class bases (#13766)

Summary
---------

PEP 695 Generics introduce a scope inside a class statement's arguments
and keywords.

```
class C[T](A[T]):  # the T in A[T] is not from the global scope but from a type-param-specfic scope
   ...
```

When doing inference on the class bases, we currently have been doing
base class expression lookups in the global scope. Not an issue without
generics (since a scope is only created when generics are present).

This change instead makes sure to stop the global scope inference from
going into expressions within this sub-scope. Since there is a separate
scope, `check_file` and friends will trigger inference on these
expressions still.

Another change as a part of this is making sure that `ClassType` looks
up its bases in the right scope.

Test Plan
----------
`cargo test --package red_knot_python_semantic generics` will run the
markdown test that previously would panic due to scope lookup issues

---------

Co-authored-by: Micha Reiser <[email protected]>
Co-authored-by: Carl Meyer <[email protected]>

crates/red_knot_python_semantic/resources/mdtest/generics.md

@@ -0,0 +1,57 @@
+# PEP 695 Generics
+
+## Class Declarations
+
+Basic PEP 695 generics
+
+```py
+class MyBox[T]:
+  data: T
+  box_model_number = 695
+  def __init__(self, data: T):
+    self.data = data
+
+# TODO not error (should be subscriptable)
+box: MyBox[int] = MyBox(5)  # error: [non-subscriptable]
+# TODO error differently (str and int don't unify)
+wrong_innards: MyBox[int] = MyBox("five")  # error: [non-subscriptable]
+# TODO reveal int
+reveal_type(box.data)  # revealed: @Todo
+
+reveal_type(MyBox.box_model_number)  # revealed: Literal[695]
+```
+
+## Subclassing
+
+```py
+class MyBox[T]:
+  data: T
+
+  def __init__(self, data: T):
+    self.data = data
+
+# TODO not error on the subscripting
+class MySecureBox[T](MyBox[T]):  # error: [non-subscriptable]
+  pass
+
+secure_box: MySecureBox[int] = MySecureBox(5)
+reveal_type(secure_box)  # revealed: MySecureBox
+# TODO reveal int
+reveal_type(secure_box.data)  # revealed: @Todo
+```
+
+## Cyclical class definition
+
+In type stubs, classes can reference themselves in their base class definitions. For example, in `typeshed`, we have `class str(Sequence[str]): ...`.
+
+This should hold true even with generics at play.
+
+```py path=a.pyi
+class Seq[T]:
+  pass
+
+# TODO not error on the subscripting
+class S[T](Seq[S]):  # error: [non-subscriptable]
+  pass
+reveal_type(S)  # revealed: Literal[S]
+```

crates/red_knot_python_semantic/src/semantic_index/ast_ids.rs

@@ -34,7 +34,10 @@ pub(crate) struct AstIds {
 
 impl AstIds {
     fn expression_id(&self, key: impl Into<ExpressionNodeKey>) -> ScopedExpressionId {
-        self.expressions_map[&key.into()]
+        let key = &key.into();
+        *self.expressions_map.get(key).unwrap_or_else(|| {
+            panic!("Could not find expression ID for {key:?}");
+        })
     }
 
     fn use_id(&self, key: impl Into<ExpressionNodeKey>) -> ScopedUseId {

crates/red_knot_python_semantic/src/types.rs

@@ -14,7 +14,7 @@ use crate::stdlib::{
     builtins_symbol_ty, types_symbol_ty, typeshed_symbol_ty, typing_extensions_symbol_ty,
 };
 use crate::types::narrow::narrowing_constraint;
-use crate::{Db, FxOrderSet, Module};
+use crate::{Db, FxOrderSet, HasTy, Module, SemanticModel};
 
 pub(crate) use self::builder::{IntersectionBuilder, UnionBuilder};
 pub use self::diagnostic::{TypeCheckDiagnostic, TypeCheckDiagnostics};
@@ -1425,7 +1425,17 @@ impl<'db> ClassType<'db> {
         class_stmt_node
             .bases()
             .iter()
-            .map(move |base_expr| definition_expression_ty(db, definition, base_expr))
+            .map(move |base_expr: &ast::Expr| {
+                if class_stmt_node.type_params.is_some() {
+                    // when we have a specialized scope, we'll look up the inference
+                    // within that scope
+                    let model: SemanticModel<'db> = SemanticModel::new(db, definition.file(db));
+                    base_expr.ty(&model)
+                } else {
+                    // Otherwise, we can do the lookup based on the definition scope
+                    definition_expression_ty(db, definition, base_expr)
+                }
+            })
     }
 
     /// Returns the class member of this class named `name`.

crates/red_knot_python_semantic/src/types/infer.rs

@@ -867,7 +867,7 @@ impl<'db> TypeInferenceBuilder<'db> {
         let ast::StmtClassDef {
             range: _,
             name,
-            type_params: _,
+            type_params,
             decorator_list,
             arguments: _,
             body: _,
@@ -885,6 +885,7 @@ impl<'db> TypeInferenceBuilder<'db> {
         let maybe_known_class = file_to_module(self.db, body_scope.file(self.db))
             .as_ref()
             .and_then(|module| KnownClass::maybe_from_module(module, name.as_str()));
+
         let class_ty = Type::Class(ClassType::new(
             self.db,
             name.id.clone(),
@@ -895,17 +896,21 @@ impl<'db> TypeInferenceBuilder<'db> {
 
         self.add_declaration_with_binding(class.into(), definition, class_ty, class_ty);
 
-        for keyword in class.keywords() {
-            self.infer_expression(&keyword.value);
-        }
+        // if there are type parameters, then the keywords and bases are within that scope
+        // and we don't need to run inference here
+        if type_params.is_none() {
+            for keyword in class.keywords() {
+                self.infer_expression(&keyword.value);
+            }
 
-        // Inference of bases deferred in stubs
-        // TODO also defer stringified generic type parameters
-        if self.are_all_types_deferred() {
-            self.types.has_deferred = true;
-        } else {
-            for base in class.bases() {
-                self.infer_expression(base);
+            // Inference of bases deferred in stubs
+            // TODO also defer stringified generic type parameters
+            if self.are_all_types_deferred() {
+                self.types.has_deferred = true;
+            } else {
+                for base in class.bases() {
+                    self.infer_expression(base);
+                }
             }
         }
     }