From 084e5464fb785f3f8fcdc54505b954eae6886906 Mon Sep 17 00:00:00 2001 From: Carl Meyer Date: Wed, 5 Jun 2024 15:13:10 -0600 Subject: [PATCH] [red-knot] support walrus expressions in type inference (#11762) ## Summary Add support for walrus expressions, both in expression type inference and in symbol definition type inference. ## Test Plan Added test. --- crates/red_knot/src/semantic.rs | 52 ++++++++++++++------- crates/red_knot/src/semantic/definitions.rs | 1 + crates/red_knot/src/semantic/types/infer.rs | 27 ++++++++++- 3 files changed, 61 insertions(+), 19 deletions(-) diff --git a/crates/red_knot/src/semantic.rs b/crates/red_knot/src/semantic.rs index 657d7b9caa..bf785fee32 100644 --- a/crates/red_knot/src/semantic.rs +++ b/crates/red_knot/src/semantic.rs @@ -240,21 +240,6 @@ impl SemanticIndexer { impl PreorderVisitor<'_> for SemanticIndexer { fn visit_expr(&mut self, expr: &ast::Expr) { - if let ast::Expr::Name(ast::ExprName { id, ctx, .. }) = expr { - let flags = match ctx { - ast::ExprContext::Load => SymbolFlags::IS_USED, - ast::ExprContext::Store => SymbolFlags::IS_DEFINED, - ast::ExprContext::Del => SymbolFlags::IS_DEFINED, - ast::ExprContext::Invalid => SymbolFlags::empty(), - }; - self.add_or_update_symbol(id, flags); - if flags.contains(SymbolFlags::IS_DEFINED) { - if let Some(curdef) = self.current_definition.clone() { - self.add_or_update_symbol_with_def(id, curdef); - } - } - } - let expression_id = self .flow_graph_builder .record_expr(self.current_flow_node()); @@ -267,7 +252,36 @@ impl PreorderVisitor<'_> for SemanticIndexer { self.expressions .insert(NodeKey::from_node(expr.into()), expression_id); - ast::visitor::preorder::walk_expr(self, expr); + + match expr { + ast::Expr::Name(ast::ExprName { id, ctx, .. }) => { + let flags = match ctx { + ast::ExprContext::Load => SymbolFlags::IS_USED, + ast::ExprContext::Store => SymbolFlags::IS_DEFINED, + ast::ExprContext::Del => SymbolFlags::IS_DEFINED, + ast::ExprContext::Invalid => SymbolFlags::empty(), + }; + self.add_or_update_symbol(id, flags); + if flags.contains(SymbolFlags::IS_DEFINED) { + if let Some(curdef) = self.current_definition.clone() { + self.add_or_update_symbol_with_def(id, curdef); + } + } + ast::visitor::preorder::walk_expr(self, expr); + } + ast::Expr::Named(node) => { + debug_assert!(self.current_definition.is_none()); + self.current_definition = + Some(Definition::NamedExpr(TypedNodeKey::from_node(node))); + // TODO walrus in comprehensions is implicitly nonlocal + self.visit_expr(&node.target); + self.current_definition = None; + self.visit_expr(&node.value); + } + _ => { + ast::visitor::preorder::walk_expr(self, expr); + } + } } fn visit_stmt(&mut self, stmt: &ast::Stmt) { @@ -394,8 +408,12 @@ impl PreorderVisitor<'_> for SemanticIndexer { debug_assert!(self.current_definition.is_none()); self.current_definition = Some(Definition::Assignment(TypedNodeKey::from_node(node))); - ast::visitor::preorder::walk_stmt(self, stmt); + for expr in &node.targets { + self.visit_expr(expr); + } + self.current_definition = None; + self.visit_expr(&node.value); } ast::Stmt::If(node) => { // we visit the if "test" condition first regardless diff --git a/crates/red_knot/src/semantic/definitions.rs b/crates/red_knot/src/semantic/definitions.rs index 0cb96b18cb..b1bd7a3ca2 100644 --- a/crates/red_knot/src/semantic/definitions.rs +++ b/crates/red_knot/src/semantic/definitions.rs @@ -19,6 +19,7 @@ pub enum Definition { FunctionDef(TypedNodeKey), Assignment(TypedNodeKey), AnnotatedAssignment(TypedNodeKey), + NamedExpr(TypedNodeKey), /// represents the implicit initial definition of every name as "unbound" Unbound, // TODO with statements, except handlers, function args... diff --git a/crates/red_knot/src/semantic/types/infer.rs b/crates/red_knot/src/semantic/types/infer.rs index 7e8c9c4139..6a11909c3f 100644 --- a/crates/red_knot/src/semantic/types/infer.rs +++ b/crates/red_knot/src/semantic/types/infer.rs @@ -161,7 +161,7 @@ pub fn infer_definition_type( let parsed = parse(db.upcast(), file_id)?; let ast = parsed.syntax(); let node = node_key.resolve_unwrap(ast.as_any_node_ref()); - // TODO handle unpacking assignment correctly (here and for AnnotatedAssignment case, below) + // TODO handle unpacking assignment infer_expr_type(db, file_id, &node.value) } Definition::AnnotatedAssignment(node_key) => { @@ -172,9 +172,15 @@ pub fn infer_definition_type( let Some(value) = &node.value else { return Ok(Type::Unknown); }; - // TODO handle unpacking assignment correctly (here and for Assignment case, above) + // TODO handle unpacking assignment infer_expr_type(db, file_id, value) } + Definition::NamedExpr(node_key) => { + let parsed = parse(db.upcast(), file_id)?; + let ast = parsed.syntax(); + let node = node_key.resolve_unwrap(ast.as_any_node_ref()); + infer_expr_type(db, file_id, &node.value) + } } } @@ -220,6 +226,7 @@ fn infer_expr_type(db: &dyn SemanticDb, file_id: FileId, expr: &ast::Expr) -> Qu // TODO add reverse bin op support if right <: left left_ty.resolve_bin_op(db, *op, right_ty) } + ast::Expr::Named(ast::ExprNamed { value, .. }) => infer_expr_type(db, file_id, value), _ => todo!("expression type resolution for {:?}", expr), } } @@ -519,4 +526,20 @@ mod tests { assert_public_type(&case, "a", "d", "Literal[-1]")?; assert_public_type(&case, "a", "e", "Literal[2]") } + + #[test] + fn walrus() -> anyhow::Result<()> { + let case = create_test()?; + + write_to_path( + &case, + "a.py", + " + x = (y := 1) + 1 + ", + )?; + + assert_public_type(&case, "a", "x", "Literal[2]")?; + assert_public_type(&case, "a", "y", "Literal[1]") + } }