[red-knot] Use the right scope when considering class bases

crates/red_knot_python_semantic/resources/mdtest/generics.md

@@ -0,0 +1,41 @@
+# 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
+```

crates/red_knot_python_semantic/src/semantic_index/ast_ids.rs

@@ -34,7 +34,13 @@ pub(crate) struct AstIds {
 
 impl AstIds {
     fn expression_id(&self, key: impl Into<ExpressionNodeKey>) -> ScopedExpressionId {
-        self.expressions_map[&key.into()]
+        let key = &key.into();
+        match self.expressions_map.get(key) {
+            Some(result) => *result,
+            None => {
+                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

@@ -1,11 +1,11 @@
-use infer::TypeInferenceBuilder;
+use infer::{TypeInference, TypeInferenceBuilder};
 use ruff_db::files::File;
 use ruff_python_ast as ast;
 
 use crate::module_resolver::file_to_module;
 use crate::semantic_index::ast_ids::HasScopedAstId;
 use crate::semantic_index::definition::{Definition, DefinitionKind};
-use crate::semantic_index::symbol::{ScopeId, ScopedSymbolId};
+use crate::semantic_index::symbol::{NodeWithScopeRef, ScopeId, ScopedSymbolId};
 use crate::semantic_index::{
     global_scope, semantic_index, symbol_table, use_def_map, BindingWithConstraints,
     BindingWithConstraintsIterator, DeclarationsIterator,
@@ -1416,10 +1416,31 @@ impl<'db> ClassType<'db> {
         let DefinitionKind::Class(class_stmt_node) = definition.kind(db) else {
             panic!("Class type definition must have DefinitionKind::Class");
         };
-        class_stmt_node
-            .bases()
-            .iter()
-            .map(move |base_expr| definition_expression_ty(db, definition, base_expr))
+
+        let has_type_params = class_stmt_node.type_params.is_some();
+
+        let type_scope_info: Option<(ScopeId<'db>, &TypeInference<'db>)> = if has_type_params {
+            let file = definition.file(db);
+            let index = semantic_index(db, file);
+            // we need to use the type param'd scope
+            let type_param_scope = index
+                .node_scope(NodeWithScopeRef::ClassTypeParameters(class_stmt_node))
+                .to_scope_id(db, file);
+            Some((type_param_scope, infer_scope_types(db, type_param_scope)))
+        } else {
+            None
+        };
+
+        let mapper = move |base_expr: &ast::Expr| {
+            if has_type_params {
+                // Safety: we calculated the inference if has_type_params
+                let (type_param_scope, inferences) = type_scope_info.unwrap();
+                inferences.expression_ty(base_expr.scoped_ast_id(db, type_param_scope))
+            } else {
+                definition_expression_ty(db, definition, base_expr)
+            }
+        };
+        class_stmt_node.bases().iter().map(mapper)
     }
 
     /// Returns the class member of this class named `name`.

crates/red_knot_python_semantic/src/types/infer.rs

@@ -853,7 +853,7 @@ impl<'db> TypeInferenceBuilder<'db> {
         let ast::StmtClassDef {
             range: _,
             name,
-            type_params: _,
+            type_params,
             decorator_list,
             arguments: _,
             body: _,
@@ -881,17 +881,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);
+                }
             }
         }
     }