[ty] Support walrus operator scope leaking from comprehensions

Author: charliermarsh

crates/ty_python_semantic/resources/mdtest/assignment/walrus.md

@@ -15,3 +15,104 @@ x = 0
 (x := x + 1)
 reveal_type(x)  # revealed: Literal[1]
 ```
+
+## Walrus in comprehensions
+
+PEP 572: Named expressions in comprehensions bind the target in the first enclosing scope that is
+not a comprehension.
+
+### List comprehension element
+
+```py
+class Iterator:
+    def __next__(self) -> int:
+        return 42
+
+class Iterable:
+    def __iter__(self) -> Iterator:
+        return Iterator()
+
+[(a := b * 2) for b in Iterable()]
+reveal_type(a)  # revealed: int
+```
+
+### Comprehension filter
+
+```py
+class Iterator:
+    def __next__(self) -> int:
+        return 42
+
+class Iterable:
+    def __iter__(self) -> Iterator:
+        return Iterator()
+
+[c for d in Iterable() if (c := d - 10) > 0]
+reveal_type(c)  # revealed: int
+```
+
+### Dict comprehension
+
+```py
+class Iterator:
+    def __next__(self) -> int:
+        return 42
+
+class Iterable:
+    def __iter__(self) -> Iterator:
+        return Iterator()
+
+{(e := f * 2): (g := f * 3) for f in Iterable()}
+reveal_type(e)  # revealed: int
+reveal_type(g)  # revealed: int
+```
+
+### Generator expression
+
+```py
+class Iterator:
+    def __next__(self) -> int:
+        return 42
+
+class Iterable:
+    def __iter__(self) -> Iterator:
+        return Iterator()
+
+list(((h := i * 2) for i in Iterable()))
+reveal_type(h)  # revealed: int
+```
+
+### Nested comprehension
+
+```py
+[[(x := y) for y in range(3)] for _ in range(3)]
+reveal_type(x)  # revealed: int
+```
+
+### Honoring `global` declaration
+
+PEP 572: the walrus honors a `global` declaration in the enclosing scope.
+
+```py
+x: int = 0
+
+def f() -> None:
+    global x
+    [(x := y) for y in range(3)]
+    reveal_type(x)  # revealed: int
+```
+
+### Honoring `nonlocal` declaration
+
+PEP 572: the walrus honors a `nonlocal` declaration in the enclosing scope.
+
+```py
+def outer() -> None:
+    x = "hello"
+
+    def inner() -> None:
+        nonlocal x
+        [(x := y) for y in range(3)]
+        reveal_type(x)  # revealed: int
+    inner()
+```

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

@@ -491,22 +491,13 @@ reveal_type(s)  # revealed: Unknown
 # error: [unresolved-reference]
 reveal_type(t)  # revealed: Unknown
 
-# TODO: these should all reveal `Unknown | int` and should not emit errors.
-# (We don't generally model elsewhere in ty that bindings from walruses
-# "leak" from comprehension scopes into outer scopes, but we should.)
-# See https://github.com/astral-sh/ruff/issues/16954
-# error: [unresolved-reference]
-reveal_type(g)  # revealed: Unknown
-# error: [unresolved-reference]
-reveal_type(i)  # revealed: Unknown
-# error: [unresolved-reference]
-reveal_type(k)  # revealed: Unknown
-# error: [unresolved-reference]
-reveal_type(m)  # revealed: Unknown
-# error: [unresolved-reference]
-reveal_type(o)  # revealed: Unknown
-# error: [unresolved-reference]
-reveal_type(q)  # revealed: Unknown
+# PEP 572: walrus targets in comprehensions leak into the enclosing scope.
+reveal_type(g)  # revealed: int
+reveal_type(i)  # revealed: int
+reveal_type(k)  # revealed: int
+reveal_type(m)  # revealed: int
+reveal_type(o)  # revealed: int
+reveal_type(q)  # revealed: int
 ```
 
 ### An annotation without a value is a definition in a stub but not a `.py` file

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -251,6 +251,23 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         false
     }
 
+    /// Returns the enclosing non-comprehension scope for walrus operator targets,
+    /// per [PEP 572]. Named expressions in comprehensions bind in the first
+    /// enclosing scope that is *not* a comprehension.
+    ///
+    /// Returns `None` if the current scope is not a comprehension.
+    ///
+    /// [PEP 572]: https://peps.python.org/pep-0572/#scope-of-the-target
+    fn enclosing_scope_for_walrus(&self) -> Option<FileScopeId> {
+        if self.scopes[self.current_scope()].kind() != ScopeKind::Comprehension {
+            return None;
+        }
+        self.scope_stack.iter().rev().skip(1).find_map(|info| {
+            (self.scopes[info.file_scope_id].kind() != ScopeKind::Comprehension)
+                .then_some(info.file_scope_id)
+        })
+    }
+
     /// Push a new loop, returning the outer loop, if any.
     fn push_loop(&mut self) -> Option<Loop> {
         self.current_scope_info_mut()
@@ -668,8 +685,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         definition
     }
 
-    fn delete_associated_bindings(&mut self, place: ScopedPlaceId) {
-        let scope = self.current_scope();
+    fn delete_associated_bindings_in_scope(&mut self, scope: FileScopeId, place: ScopedPlaceId) {
         // Don't delete associated bindings if the scope is a class scope & place is a name (it's never visible to nested scopes)
         if self.scopes[scope].kind() == ScopeKind::Class && place.is_symbol() {
             return;
@@ -687,19 +703,18 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         self.current_use_def_map_mut().delete_binding(place);
     }
 
-    /// Push a new [`Definition`] onto the list of definitions
-    /// associated with the `definition_node` AST node.
+    /// Create a [`Definition`] for `definition_node` in the given `scope`, recording it in
+    /// that scope's place table and use-def map.
     ///
-    /// Returns a 2-element tuple, where the first element is the newly created [`Definition`]
-    /// and the second element is the number of definitions that are now associated with
-    /// `definition_node`.
+    /// Returns a 2-element tuple: the newly created [`Definition`] and the total number of
+    /// definitions now associated with the AST node.
     ///
-    /// This method should only be used when adding a definition associated with a `*` import.
-    /// All other nodes can only ever be associated with exactly 1 or 0 [`Definition`]s.
-    /// For any node other than an [`ast::Alias`] representing a `*` import,
-    /// prefer to use `self.add_definition()`, which ensures that this invariant is maintained.
-    fn push_additional_definition(
+    /// This is the low-level helper; callers that add definitions to the **current** scope
+    /// should normally use [`Self::push_additional_definition`] or [`Self::add_definition`],
+    /// which also update lazy snapshots and the try-node context stack.
+    fn add_definition_in_scope(
         &mut self,
+        scope: FileScopeId,
         place: ScopedPlaceId,
         definition_node: impl Into<DefinitionNodeRef<'ast, 'db>>,
     ) -> (Definition<'db>, usize) {
@@ -711,14 +726,8 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         let category = kind.category(self.source_type.is_stub(), self.module);
         let is_reexported = kind.is_reexported();
 
-        let definition: Definition<'db> = Definition::new(
-            self.db,
-            self.file,
-            self.current_scope(),
-            place,
-            kind,
-            is_reexported,
-        );
+        let definition: Definition<'db> =
+            Definition::new(self.db, self.file, scope, place, kind, is_reexported);
 
         let num_definitions = {
             let definitions = self.add_entry_for_definition_key(definition_node.key());
@@ -727,25 +736,53 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         };
 
         if category.is_binding() {
-            self.mark_place_bound(place);
+            self.place_tables[scope].mark_bound(place);
         }
         if category.is_declaration() {
-            self.mark_place_declared(place);
+            self.place_tables[scope].mark_declared(place);
         }
 
-        let use_def = self.current_use_def_map_mut();
         match category {
             DefinitionCategory::DeclarationAndBinding => {
-                use_def.record_declaration_and_binding(place, definition);
-                self.delete_associated_bindings(place);
+                self.use_def_maps[scope].record_declaration_and_binding(place, definition);
+                self.delete_associated_bindings_in_scope(scope, place);
+            }
+            DefinitionCategory::Declaration => {
+                self.use_def_maps[scope].record_declaration(place, definition);
             }
-            DefinitionCategory::Declaration => use_def.record_declaration(place, definition),
             DefinitionCategory::Binding => {
-                use_def.record_binding(place, definition);
-                self.delete_associated_bindings(place);
+                self.use_def_maps[scope].record_binding(place, definition);
+                self.delete_associated_bindings_in_scope(scope, place);
             }
         }
 
+        (definition, num_definitions)
+    }
+
+    /// Push a new [`Definition`] onto the list of definitions
+    /// associated with the `definition_node` AST node in the **current** scope.
+    ///
+    /// Returns a 2-element tuple, where the first element is the newly created [`Definition`]
+    /// and the second element is the number of definitions that are now associated with
+    /// `definition_node`.
+    ///
+    /// This method should only be used when adding a definition associated with a `*` import.
+    /// All other nodes can only ever be associated with exactly 1 or 0 [`Definition`]s.
+    /// For any node other than an [`ast::Alias`] representing a `*` import,
+    /// prefer to use `self.add_definition()`, which ensures that this invariant is maintained.
+    fn push_additional_definition(
+        &mut self,
+        place: ScopedPlaceId,
+        definition_node: impl Into<DefinitionNodeRef<'ast, 'db>>,
+    ) -> (Definition<'db>, usize) {
+        let scope = self.current_scope();
+        let (definition, num_definitions) =
+            self.add_definition_in_scope(scope, place, definition_node);
+
+        let category = definition
+            .kind(self.db)
+            .category(self.source_type.is_stub(), self.module);
+
         if category.is_binding() {
             if let Some(id) = place.as_symbol() {
                 self.update_lazy_snapshots(id);
@@ -2632,10 +2669,28 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                                 );
                             }
                             Some(CurrentAssignment::Named(named)) => {
-                                // TODO(dhruvmanila): If the current scope is a comprehension, then the
-                                // named expression is implicitly nonlocal. This is yet to be
-                                // implemented.
-                                self.add_definition(place_id, named);
+                                if let Some(enclosing_scope) = self.enclosing_scope_for_walrus() {
+                                    // PEP 572: walrus in comprehension binds in enclosing scope.
+                                    let target_name = named
+                                        .target
+                                        .as_name_expr()
+                                        .expect("target should be a Name expression")
+                                        .id
+                                        .clone();
+                                    let (symbol_id, added) = self.place_tables[enclosing_scope]
+                                        .add_symbol(Symbol::new(target_name));
+                                    if added {
+                                        self.use_def_maps[enclosing_scope]
+                                            .add_place(symbol_id.into());
+                                    }
+                                    self.add_definition_in_scope(
+                                        enclosing_scope,
+                                        symbol_id.into(),
+                                        named,
+                                    );
+                                } else {
+                                    self.add_definition(place_id, named);
+                                }
                             }
                             Some(CurrentAssignment::Comprehension {
                                 unpack,
@@ -2690,11 +2745,16 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                 walk_expr(self, expr);
             }
             ast::Expr::Named(node) => {
-                // TODO walrus in comprehensions is implicitly nonlocal
                 self.visit_expr(&node.value);
 
                 // See https://peps.python.org/pep-0572/#differences-between-assignment-expressions-and-assignment-statements
                 if node.target.is_name_expr() {
+                    // PEP 572: walrus in comprehension binds in the enclosing scope.
+                    // Make the value a standalone expression so inference can evaluate
+                    // it in the comprehension scope where the iteration variables are visible.
+                    if self.enclosing_scope_for_walrus().is_some() {
+                        self.add_standalone_expression(&node.value);
+                    }
                     self.push_assignment(node.into());
                     self.visit_expr(&node.target);
                     self.pop_assignment();

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -11290,10 +11290,24 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
     fn infer_named_expression(&mut self, named: &ast::ExprNamed) -> Type<'db> {
         // See https://peps.python.org/pep-0572/#differences-between-assignment-expressions-and-assignment-statements
         if named.target.is_name_expr() {
-            let definition = self.index.expect_single_definition(named);
-            let result = infer_definition_types(self.db(), definition);
-            self.extend_definition(result);
-            result.binding_type(definition)
+            let db = self.db();
+            let file_scope_id = self.scope().file_scope_id(db);
+
+            if self.index.scope(file_scope_id).kind() == ScopeKind::Comprehension {
+                // PEP 572: walrus in comprehension binds in the enclosing scope.
+                // Infer the value via its standalone expression in this scope;
+                // the definition lives in the enclosing scope and will be
+                // inferred when that scope needs it.
+                let expression = self.index.expression(named.value.as_ref());
+                let result = infer_expression_types(db, expression, TypeContext::default());
+                self.extend_expression(result);
+                result.expression_type(named.value.as_ref())
+            } else {
+                let definition = self.index.expect_single_definition(named);
+                let result = infer_definition_types(db, definition);
+                self.extend_definition(result);
+                result.binding_type(definition)
+            }
         } else {
             // For syntactically invalid targets, we still need to run type inference:
             self.infer_expression(&named.target, TypeContext::default());
@@ -11316,7 +11330,17 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
 
         let add = self.add_binding(named.target.as_ref().into(), definition);
 
-        let ty = self.infer_expression(value, add.type_context());
+        // PEP 572: walrus in a comprehension binds in the enclosing scope, but
+        // the value references comprehension-scoped variables. The builder
+        // registers the value as a standalone expression in the comprehension
+        // scope so we can infer it there. We must not `extend` the result
+        // because expression IDs are only meaningful within their own scope.
+        let ty = if let Some(expression) = self.index.try_expression(value.as_ref()) {
+            let result = infer_expression_types(self.db(), expression, add.type_context());
+            result.expression_type(value.as_ref())
+        } else {
+            self.infer_expression(value, add.type_context())
+        };
         self.store_expression_type(target, ty);
         add.insert(self, ty)
     }