[ty] Support walrus operator scope leaking from comprehensions

Lazy

Avoid clone

Author: charliermarsh

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

@@ -15,3 +15,153 @@ 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
+```
+
+### Class body comprehension
+
+```py
+class C:
+    [(x := y) for y in range(3)]
+    reveal_type(x)  # revealed: int
+```
+
+### First generator `iter`
+
+The `iter` of the first generator is evaluated in the enclosing scope. A walrus there should bind in
+the enclosing scope as usual (no comprehension scope is involved).
+
+```py
+def returns_list() -> list[int]:
+    return [1, 2, 3]
+
+[x for x in (y := returns_list())]
+reveal_type(y)  # revealed: list[int]
+```
+
+### Nested comprehension
+
+```py
+[[(x := y) for y in range(3)] for _ in range(3)]
+reveal_type(x)  # revealed: int
+```
+
+### Updates lazy snapshots in nested scopes
+
+```py
+def returns_str() -> str:
+    return "foo"
+
+def outer() -> None:
+    x = returns_str()
+
+    def inner() -> None:
+        reveal_type(x)  # revealed: str | int
+    [(x := y) for y in range(3)]
+    inner()
+```
+
+### Possibly defined in `except` handlers
+
+```py
+def could_raise() -> list[int]:
+    return [1]
+
+try:
+    [(y := n) for n in could_raise()]
+except:
+    # error: [possibly-unresolved-reference]
+    reveal_type(y)  # 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

@@ -285,6 +285,28 @@ 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, usize)> {
+        if self.scopes[self.current_scope()].kind() != ScopeKind::Comprehension {
+            return None;
+        }
+        self.scope_stack
+            .iter()
+            .enumerate()
+            .rev()
+            .skip(1)
+            .find_map(|(index, info)| {
+                (self.scopes[info.file_scope_id].kind() != ScopeKind::Comprehension)
+                    .then_some((info.file_scope_id, index))
+            })
+    }
+
     /// Push a new loop, returning the outer loop, if any.
     fn push_loop(&mut self) -> Option<Loop> {
         self.current_scope_info_mut()
@@ -801,10 +823,28 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, '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, scope_index)) = 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.push_additional_definition_in_scope(
+                        enclosing_scope,
+                        scope_index,
+                        symbol_id.into(),
+                        named,
+                    );
+                } else {
+                    self.add_definition(place_id, named);
+                }
             }
             Some(CurrentAssignment::Comprehension {
                 unpack,
@@ -3167,11 +3207,16 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                 }
             }
             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

@@ -6209,10 +6209,23 @@ 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();
+
+            if self.scope().node(db).scope_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());
@@ -6235,7 +6248,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)
     }