further improve handling of walruses in lambdas and comprehensions

Author: AlexWaygood

crates/red_knot_python_semantic/resources/mdtest/import/star.md

@@ -212,9 +212,9 @@ print((
 ))
 ```
 
-### Definitions in comprehension-like scopes are not global definitions
+### Definitions in function-like scopes are not global definitions
 
-Except for some cases involving walrus expressions.
+Except for some cases involving walrus expressions inside comprehension scopes.
 
 `a.py`:
 
@@ -231,10 +231,13 @@ class Iterable:
 {b for b in Iterable()}
 {c: c for c in Iterable()}
 (d for d in Iterable())
+lambda e: (f := 42)
 
-[(e := f * 2) for f in Iterable()]
-[g for h in Iterable() if (g := h - 10) > 0]
-{(i := j * 2): (k := j * 3) for j in Iterable()}
+# Definitions created by walruses in a comprehension scope are unique;
+# they "leak out" of the scope and are stored in the surrounding scope
+[(g := h * 2) for h in Iterable()]
+[i for j in Iterable() if (i := j - 10) > 0]
+{(k := l * 2): (m := l * 3) for l in Iterable()}
 ```
 
 `b.py`:
@@ -251,6 +254,8 @@ reveal_type(c)  # revealed: Unknown
 # error: [unresolved-reference]
 reveal_type(d)  # revealed: Unknown
 # error: [unresolved-reference]
+reveal_type(e)  # revealed: Unknown
+# error: [unresolved-reference]
 reveal_type(f)  # revealed: Unknown
 # error: [unresolved-reference]
 reveal_type(h)  # revealed: Unknown
@@ -260,10 +265,10 @@ reveal_type(j)  # revealed: Unknown
 # TODO: these should all reveal `int`
 # (we don't generally model elsewhere in red-knot that bindings from walruses
 # "leak" from comprehension scopes into outer scopes, but we should)
-reveal_type(e)  # revealed: Unknown
 reveal_type(g)  # revealed: Unknown
 reveal_type(i)  # revealed: Unknown
 reveal_type(k)  # revealed: Unknown
+reveal_type(m)  # revealed: Unknown
 ```
 
 ### An annotation without a value is a definition in a stub but not a `.py` file

crates/red_knot_python_semantic/src/semantic_index/re_exports.rs

@@ -221,48 +221,46 @@ impl<'db> Visitor<'db> for ExportFinder<'db> {
                 }
             }
 
-            ast::Expr::SetComp(ast::ExprSetComp {
-                elt,
-                generators,
-                range: _,
-            })
-            | ast::Expr::ListComp(ast::ExprListComp {
-                elt,
-                generators,
-                range: _,
-            })
-            | ast::Expr::Generator(ast::ExprGenerator {
-                elt,
-                generators,
-                range: _,
-                parenthesized: _,
-            }) => {
-                let mut walrus_finder = WalrusFinder {
-                    export_finder: self,
-                };
-                walrus_finder.visit_expr(elt);
-                for generator in generators {
-                    walrus_finder.visit_comprehension(generator);
-                }
-            }
+            ast::Expr::Lambda(_)
+            | ast::Expr::BooleanLiteral(_)
+            | ast::Expr::NoneLiteral(_)
+            | ast::Expr::NumberLiteral(_)
+            | ast::Expr::BytesLiteral(_)
+            | ast::Expr::EllipsisLiteral(_)
+            | ast::Expr::StringLiteral(_) => {}
 
-            ast::Expr::DictComp(ast::ExprDictComp {
-                key,
-                value,
-                generators,
-                range: _,
-            }) => {
+            // Walrus definitions "leak" from comprehension scopes into the comprehension's
+            // enclosing scope; they thus need special handling
+            ast::Expr::SetComp(_)
+            | ast::Expr::ListComp(_)
+            | ast::Expr::Generator(_)
+            | ast::Expr::DictComp(_) => {
                 let mut walrus_finder = WalrusFinder {
                     export_finder: self,
                 };
-                walrus_finder.visit_expr(key);
-                walrus_finder.visit_expr(value);
-                for generator in generators {
-                    walrus_finder.visit_comprehension(generator);
-                }
+                walk_expr(&mut walrus_finder, expr);
             }
 
-            _ => walk_expr(self, expr),
+            ast::Expr::BoolOp(_)
+            | ast::Expr::Named(_)
+            | ast::Expr::BinOp(_)
+            | ast::Expr::UnaryOp(_)
+            | ast::Expr::If(_)
+            | ast::Expr::Attribute(_)
+            | ast::Expr::Subscript(_)
+            | ast::Expr::Starred(_)
+            | ast::Expr::Call(_)
+            | ast::Expr::Compare(_)
+            | ast::Expr::Yield(_)
+            | ast::Expr::YieldFrom(_)
+            | ast::Expr::FString(_)
+            | ast::Expr::Tuple(_)
+            | ast::Expr::List(_)
+            | ast::Expr::Slice(_)
+            | ast::Expr::IpyEscapeCommand(_)
+            | ast::Expr::Dict(_)
+            | ast::Expr::Set(_)
+            | ast::Expr::Await(_) => walk_expr(self, expr),
         }
     }
 }
@@ -274,10 +272,20 @@ struct WalrusFinder<'a, 'db> {
 impl<'db> Visitor<'db> for WalrusFinder<'_, 'db> {
     fn visit_expr(&mut self, expr: &'db ast::Expr) {
         match expr {
+            // anything that creates a nested scope or that cannot contain a walrus
+            // can be short-circuited
             ast::Expr::DictComp(_)
             | ast::Expr::SetComp(_)
             | ast::Expr::ListComp(_)
-            | ast::Expr::Generator(_) => {}
+            | ast::Expr::Generator(_)
+            | ast::Expr::Lambda(_)
+            | ast::Expr::BooleanLiteral(_)
+            | ast::Expr::NoneLiteral(_)
+            | ast::Expr::NumberLiteral(_)
+            | ast::Expr::BytesLiteral(_)
+            | ast::Expr::EllipsisLiteral(_)
+            | ast::Expr::StringLiteral(_)
+            | ast::Expr::Name(_) => {}
 
             ast::Expr::Named(ast::ExprNamed {
                 target,
@@ -294,7 +302,25 @@ impl<'db> Visitor<'db> for WalrusFinder<'_, 'db> {
                 }
             }
 
-            _ => walk_expr(self, expr),
+            ast::Expr::BoolOp(_)
+            | ast::Expr::BinOp(_)
+            | ast::Expr::UnaryOp(_)
+            | ast::Expr::If(_)
+            | ast::Expr::Attribute(_)
+            | ast::Expr::Subscript(_)
+            | ast::Expr::Starred(_)
+            | ast::Expr::Call(_)
+            | ast::Expr::Compare(_)
+            | ast::Expr::Yield(_)
+            | ast::Expr::YieldFrom(_)
+            | ast::Expr::FString(_)
+            | ast::Expr::Tuple(_)
+            | ast::Expr::List(_)
+            | ast::Expr::Slice(_)
+            | ast::Expr::IpyEscapeCommand(_)
+            | ast::Expr::Dict(_)
+            | ast::Expr::Set(_)
+            | ast::Expr::Await(_) => walk_expr(self, expr),
         }
     }
 }

crates/red_knot_python_semantic/src/semantic_model.rs

@@ -149,9 +149,8 @@ macro_rules! impl_binding_has_ty {
             #[inline]
             fn inferred_type<'db>(&self, model: &SemanticModel<'db>) -> Type<'db> {
                 let index = semantic_index(model.db, model.file);
-                let definitions = index.definitions(self);
-                debug_assert_eq!(definitions.len(), 1);
-                binding_type(model.db, definitions[0])
+                let binding = index.expect_single_definition(self);
+                binding_type(model.db, binding)
             }
         }
     };

crates/red_knot_python_semantic/src/types/infer.rs

@@ -1339,7 +1339,8 @@ impl<'db> TypeInferenceBuilder<'db> {
 
     fn infer_definition(&mut self, node: impl Into<DefinitionNodeKey> + std::fmt::Debug + Copy) {
         let definition = self.index.expect_single_definition(node);
-        self.extend(infer_definition_types(self.db(), definition));
+        let result = infer_definition_types(self.db(), definition);
+        self.extend(result);
     }
 
     fn infer_function_definition_statement(&mut self, function: &ast::StmtFunctionDef) {