[red-knot] simplify type lookup in function/class definitions (#14303)

When we look up the types of class bases or keywords (`metaclass`), we
currently do this little dance: if there are type params, then look up
the type using `SemanticModel` in the type-params scope, if not, look up
the type directly in the definition's own scope, with support for
deferred types.

With inference of function parameter types, I'm now adding another case
of this same dance, so I'm motivated to make it a bit more ergonomic.

Add support to `definition_expression_ty` to handle any sub-expression
of a definition, whether it is in the definition's own scope or in a
type-params sub-scope.

Related to both #13693 and #14161.

crates/red_knot_python_semantic/src/types.rs

@@ -20,7 +20,7 @@ use crate::symbol::{Boundness, Symbol};
 use crate::types::diagnostic::TypeCheckDiagnosticsBuilder;
 use crate::types::mro::{ClassBase, Mro, MroError, MroIterator};
 use crate::types::narrow::narrowing_constraint;
-use crate::{Db, FxOrderSet, HasTy, Module, Program, SemanticModel};
+use crate::{Db, FxOrderSet, Module, Program};
 
 pub(crate) use self::builder::{IntersectionBuilder, UnionBuilder};
 pub use self::diagnostic::{TypeCheckDiagnostic, TypeCheckDiagnostics};
@@ -191,19 +191,31 @@ fn declaration_ty<'db>(db: &'db dyn Db, definition: Definition<'db>) -> Type<'db
 
 /// Infer the type of a (possibly deferred) sub-expression of a [`Definition`].
 ///
+/// Supports expressions that are evaluated within a type-params sub-scope.
+///
 /// ## Panics
 /// If the given expression is not a sub-expression of the given [`Definition`].
 fn definition_expression_ty<'db>(
     db: &'db dyn Db,
     definition: Definition<'db>,
     expression: &ast::Expr,
 ) -> Type<'db> {
-    let expr_id = expression.scoped_ast_id(db, definition.scope(db));
-    let inference = infer_definition_types(db, definition);
-    if let Some(ty) = inference.try_expression_ty(expr_id) {
-        ty
+    let file = definition.file(db);
+    let index = semantic_index(db, file);
+    let file_scope = index.expression_scope_id(expression);
+    let scope = file_scope.to_scope_id(db, file);
+    let expr_id = expression.scoped_ast_id(db, scope);
+    if scope == definition.scope(db) {
+        // expression is in the definition scope
+        let inference = infer_definition_types(db, definition);
+        if let Some(ty) = inference.try_expression_ty(expr_id) {
+            ty
+        } else {
+            infer_deferred_types(db, definition).expression_ty(expr_id)
+        }
     } else {
-        infer_deferred_types(db, definition).expression_ty(expr_id)
+        // expression is in a type-params sub-scope
+        infer_scope_types(db, scope).expression_ty(expr_id)
     }
 }
 
@@ -2422,26 +2434,13 @@ impl<'db> Class<'db> {
     fn explicit_bases_query(self, db: &'db dyn Db) -> Box<[Type<'db>]> {
         let class_stmt = self.node(db);
 
-        if class_stmt.type_params.is_some() {
-            // when we have a specialized scope, we'll look up the inference
-            // within that scope
-            let model = SemanticModel::new(db, self.file(db));
-
-            class_stmt
-                .bases()
-                .iter()
-                .map(|base| base.ty(&model))
-                .collect()
-        } else {
-            // Otherwise, we can do the lookup based on the definition scope
-            let class_definition = semantic_index(db, self.file(db)).definition(class_stmt);
+        let class_definition = semantic_index(db, self.file(db)).definition(class_stmt);
 
-            class_stmt
-                .bases()
-                .iter()
-                .map(|base_node| definition_expression_ty(db, class_definition, base_node))
-                .collect()
-        }
+        class_stmt
+            .bases()
+            .iter()
+            .map(|base_node| definition_expression_ty(db, class_definition, base_node))
+            .collect()
     }
 
     fn file(self, db: &dyn Db) -> File {
@@ -2502,16 +2501,9 @@ impl<'db> Class<'db> {
             .as_ref()?
             .find_keyword("metaclass")?
             .value;
-        Some(if class_stmt.type_params.is_some() {
-            // when we have a specialized scope, we'll look up the inference
-            // within that scope
-            let model = SemanticModel::new(db, self.file(db));
-            metaclass_node.ty(&model)
-        } else {
-            // Otherwise, we can do the lookup based on the definition scope
-            let class_definition = semantic_index(db, self.file(db)).definition(class_stmt);
-            definition_expression_ty(db, class_definition, metaclass_node)
-        })
+        let class_definition = semantic_index(db, self.file(db)).definition(class_stmt);
+        let metaclass_ty = definition_expression_ty(db, class_definition, metaclass_node);
+        Some(metaclass_ty)
     }
 
     /// Return the metaclass of this class, or `Unknown` if the metaclass cannot be inferred.