From b9c1dd3d625ae92062a7766f0215b6ac27d1f6f1 Mon Sep 17 00:00:00 2001 From: Carl Meyer Date: Wed, 28 Aug 2024 10:53:09 -0700 Subject: [PATCH] generic lookup for maybe-deferred expr in a definition --- crates/red_knot_python_semantic/src/types.rs | 34 +++++++++++++------ .../src/types/infer.rs | 5 +++ 2 files changed, 29 insertions(+), 10 deletions(-) diff --git a/crates/red_knot_python_semantic/src/types.rs b/crates/red_knot_python_semantic/src/types.rs index b62b7c853dbe4..c9b698a13ba04 100644 --- a/crates/red_knot_python_semantic/src/types.rs +++ b/crates/red_knot_python_semantic/src/types.rs @@ -90,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 @@ -344,18 +362,14 @@ impl<'db> ClassType<'db> { /// # Panics: /// If `definition` is not a `DefinitionKind::Class`. pub fn bases(&self, db: &'db dyn Db) -> impl Iterator> { - let DefinitionKind::Class(class_stmt_node) = self.definition(db).node(db) else { - panic!("Class type definition must have DefinitionKind::Class"); - }; let definition = self.definition(db); - let inference = if definition.file(db).is_stub(db.upcast()) { - infer_deferred_types(db, definition) - } else { - infer_definition_types(db, definition) + 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| { - inference.expression_ty(base_expr.scoped_ast_id(db, definition.scope(db))) - }) + 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`. diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs index b97a35b63ceaa..c0e2df24a5334 100644 --- a/crates/red_knot_python_semantic/src/types/infer.rs +++ b/crates/red_knot_python_semantic/src/types/infer.rs @@ -177,6 +177,10 @@ impl<'db> TypeInference<'db> { self.expressions[&expression] } + pub(crate) fn try_expression_ty(&self, expression: ScopedExpressionId) -> Option> { + self.expressions.get(&expression).copied() + } + pub(crate) fn definition_ty(&self, definition: Definition<'db>) -> Type<'db> { self.definitions[&definition] } @@ -658,6 +662,7 @@ impl<'db> TypeInferenceBuilder<'db> { } // inference of bases deferred in stubs + // TODO also defer stringified generic type parameters if !self.is_stub() { for base in class.bases() { self.infer_expression(base);