[red-knot] support deferred evaluation of type expressions

crates/red_knot_python_semantic/src/types.rs

@@ -1,8 +1,9 @@
 use ruff_db::files::File;
-use ruff_python_ast::name::Name;
+use ruff_python_ast as ast;
 
 use crate::builtins::builtins_scope;
-use crate::semantic_index::definition::Definition;
+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::{
     global_scope, semantic_index, symbol_table, use_def_map, DefinitionWithConstraints,
@@ -14,7 +15,8 @@ use crate::{Db, FxOrderSet};
 pub(crate) use self::builder::{IntersectionBuilder, UnionBuilder};
 pub(crate) use self::diagnostic::TypeCheckDiagnostics;
 pub(crate) use self::infer::{
-    infer_definition_types, infer_expression_types, infer_scope_types, TypeInference,
+    infer_deferred_types, infer_definition_types, infer_expression_types, infer_scope_types,
+    TypeInference,
 };
 
 mod builder;
@@ -88,6 +90,24 @@ pub(crate) fn definition_ty<'db>(db: &'db dyn Db, definition: Definition<'db>) -
     inference.definition_ty(definition)
 }
 
+/// Infer the type of a (possibly deferred) sub-expression of a [`Definition`].
+///
+/// ## Panics
+/// If the given expression is not a sub-expression of the given [`Definition`].
+pub(crate) 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
+    } else {
+        infer_deferred_types(db, definition).expression_ty(expr_id)
+    }
+}
+
 /// Infer the combined type of an array of [`Definition`]s, plus one optional "unbound type".
 ///
 /// Will return a union if there is more than one definition, or at least one plus an unbound
@@ -243,7 +263,7 @@ impl<'db> Type<'db> {
     /// us to explicitly consider whether to handle an error or propagate
     /// it up the call stack.
     #[must_use]
-    pub fn member(&self, db: &'db dyn Db, name: &Name) -> Type<'db> {
+    pub fn member(&self, db: &'db dyn Db, name: &ast::name::Name) -> Type<'db> {
         match self {
             Type::Any => Type::Any,
             Type::Never => {
@@ -314,7 +334,7 @@ impl<'db> Type<'db> {
 #[salsa::interned]
 pub struct FunctionType<'db> {
     /// name of the function at definition
-    pub name: Name,
+    pub name: ast::name::Name,
 
     /// types of all decorators on this function
     decorators: Vec<Type<'db>>,
@@ -329,19 +349,33 @@ impl<'db> FunctionType<'db> {
 #[salsa::interned]
 pub struct ClassType<'db> {
     /// Name of the class at definition
-    pub name: Name,
+    pub name: ast::name::Name,
 
-    /// Types of all class bases
-    bases: Vec<Type<'db>>,
+    definition: Definition<'db>,
 
     body_scope: ScopeId<'db>,
 }
 
 impl<'db> ClassType<'db> {
+    /// Return an iterator over the types of this class's bases.
+    ///
+    /// # Panics:
+    /// If `definition` is not a `DefinitionKind::Class`.
+    pub fn bases(&self, db: &'db dyn Db) -> impl Iterator<Item = Type<'db>> {
+        let definition = self.definition(db);
+        let DefinitionKind::Class(class_stmt_node) = definition.node(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))
+    }
+
     /// Returns the class member of this class named `name`.
     ///
     /// The member resolves to a member of the class itself or any of its bases.
-    pub fn class_member(self, db: &'db dyn Db, name: &Name) -> Type<'db> {
+    pub fn class_member(self, db: &'db dyn Db, name: &ast::name::Name) -> Type<'db> {
         let member = self.own_class_member(db, name);
         if !member.is_unbound() {
             return member;
@@ -351,12 +385,12 @@ impl<'db> ClassType<'db> {
     }
 
     /// Returns the inferred type of the class member named `name`.
-    pub fn own_class_member(self, db: &'db dyn Db, name: &Name) -> Type<'db> {
+    pub fn own_class_member(self, db: &'db dyn Db, name: &ast::name::Name) -> Type<'db> {
         let scope = self.body_scope(db);
         symbol_ty_by_name(db, scope, name)
     }
 
-    pub fn inherited_class_member(self, db: &'db dyn Db, name: &Name) -> Type<'db> {
+    pub fn inherited_class_member(self, db: &'db dyn Db, name: &ast::name::Name) -> Type<'db> {
         for base in self.bases(db) {
             let member = base.member(db, name);
             if !member.is_unbound() {