call ensure_expr before narrowing in comprehensions

Summary:
fixes #350

Previously we calculated the narrowing, and then called ensure_expr on every sub-expression using a visitor. This meant that the narrows weren't being applied properly.

We need to bind any names declared by walruses before we do the narrowing, which requires moving the ensure_expr calls into bind_comprehensions.

Reviewed By: rchen152

Differential Revision: D75472199

fbshipit-source-id: 9865df98c9cf3bc84dd21e0e2e3d01a9181ae72b

Author: yangdanny97

pyrefly/lib/binding/expr.rs

@@ -239,11 +239,13 @@ impl<'a> BindingsBuilder<'a> {
     fn bind_comprehensions(&mut self, range: TextRange, comprehensions: &mut [Comprehension]) {
         self.scopes.push(Scope::comprehension(range));
         for comp in comprehensions {
+            self.ensure_expr(&mut comp.iter);
             self.scopes.add_lvalue_to_current_static(&comp.target);
             let make_binding =
                 |ann| Binding::IterableValue(ann, comp.iter.clone(), IsAsync::new(comp.is_async));
             self.bind_target(&mut comp.target, &make_binding);
-            for x in comp.ifs.iter() {
+            for x in comp.ifs.iter_mut() {
+                self.ensure_expr(x);
                 let narrow_ops = NarrowOps::from_expr(self, Some(x));
                 self.bind_narrow_ops(&narrow_ops, comp.range);
             }
@@ -360,7 +362,7 @@ impl<'a> BindingsBuilder<'a> {
 
     /// Execute through the expr, ensuring every name has a binding.
     pub fn ensure_expr(&mut self, x: &mut Expr) {
-        let new_scope = match x {
+        match x {
             Expr::If(x) => {
                 // Ternary operation. We treat it like an if/else statement.
                 let base = self.scopes.clone_current_flow();
@@ -560,46 +562,48 @@ impl<'a> BindingsBuilder<'a> {
             {
                 // Control flow doesn't proceed after sys.exit(), exit(), quit(), or os._exit().
                 self.scopes.mark_flow_termination();
-                false
             }
             Expr::Name(x) => {
                 let name = Ast::expr_name_identifier(x.clone());
                 let binding = self
                     .lookup_name(&name.id, LookupKind::Regular)
                     .map(Binding::Forward);
                 self.ensure_name(&name, binding);
-                false
             }
             Expr::ListComp(x) => {
                 self.bind_comprehensions(x.range, &mut x.generators);
-                true
+                self.ensure_expr(&mut x.elt);
+                self.scopes.pop();
+                return;
             }
             Expr::SetComp(x) => {
                 self.bind_comprehensions(x.range, &mut x.generators);
-                true
+                self.ensure_expr(&mut x.elt);
+                self.scopes.pop();
+                return;
             }
             Expr::DictComp(x) => {
                 self.bind_comprehensions(x.range, &mut x.generators);
-                true
+                self.ensure_expr(&mut x.key);
+                self.ensure_expr(&mut x.value);
+                self.scopes.pop();
+                return;
             }
             Expr::Generator(x) => {
                 self.bind_comprehensions(x.range, &mut x.generators);
-                true
+                self.ensure_expr(&mut x.elt);
+                self.scopes.pop();
+                return;
             }
             Expr::Yield(x) => {
                 self.record_yield(x.clone());
-                false
             }
             Expr::YieldFrom(x) => {
                 self.record_yield_from(x.clone());
-                false
             }
-            _ => false,
-        };
-        x.recurse_mut(&mut |x| self.ensure_expr(x));
-        if new_scope {
-            self.scopes.pop();
+            _ => {}
         }
+        x.recurse_mut(&mut |x| self.ensure_expr(x));
     }
 
     /// Execute through the expr, ensuring every name has a binding.

pyrefly/lib/test/narrow.rs

@@ -412,6 +412,42 @@ if (x := f()) is None:
     "#,
 );
 
+testcase!(
+    test_walrus_comprehension_if_simple,
+    r#"
+from typing import assert_type
+def f(xs: list[int | None]):
+    ys = [y111 for x in xs if (y111 := x)]
+    assert_type(ys, list[int])
+    "#,
+);
+
+testcase!(
+    test_walrus_comprehension_if_function,
+    r#"
+from typing import assert_type
+def get_y(x: int | None) -> int | None:
+    return x
+def f(xs: list[int | None]):
+    ys = [y111 for x in xs if (y111 := get_y(x))]
+    assert_type(ys, list[int])
+    "#,
+);
+
+testcase!(
+    test_walrus_generator_if,
+    r#"
+from typing import Sequence
+def foo(x: int) -> int | None:
+    return (x + 5) if x % 2 else None
+foos: Sequence[int] = tuple(
+    maybe_foo
+    for x in range(10)
+    if (maybe_foo := foo(x)) is not None
+)
+    "#,
+);
+
 testcase!(
     test_match_enum_fallback,
     r#"
@@ -866,20 +902,16 @@ def f(x: Any):
 );
 
 testcase!(
-    bug = "TODO(stroxler): We are binding narrows too early and miss errors in missing attributes that get narrowed.",
     test_listcomp_if_control_flow,
     r#"
 class C: pass
 class D(C): pass
 def accepts_d(x: D) -> None: pass
 def f(x: list[C], z: C):
-    # if statement control flow works as expected, we get an error.
     if accepts_d(z) and isinstance(z, D):  # E: Argument `C` is not assignable to parameter `x` with type `D`
         pass
-    # But here we don't get an error, it appears the narrowed binding is applied too early.
-    [y for y in x if (accepts_d(y) and isinstance(y, D))]
-    # Here we don't get an error, for the same reason - we narrow too early, hiding the no-such-attribute error.
-    [None for y in x if C.error]
+    [y for y in x if (accepts_d(y) and isinstance(y, D))]  # E: Argument `C` is not assignable to parameter `x` with type `D` in function `accepts_d`
+    [None for y in x if C.error]  # E: Class `C` has no class attribute `error`
     "#,
 );