[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

@@ -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

crates/red_knot/src/semantic/definitions.rs

@@ -19,6 +19,7 @@ pub enum Definition {
     FunctionDef(TypedNodeKey<ast::StmtFunctionDef>),
     Assignment(TypedNodeKey<ast::StmtAssign>),
     AnnotatedAssignment(TypedNodeKey<ast::StmtAnnAssign>),
+    NamedExpr(TypedNodeKey<ast::ExprNamed>),
     /// represents the implicit initial definition of every name as "unbound"
     Unbound,
     // TODO with statements, except handlers, function args...

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]")
+    }
 }