fix: don't flag names as unused when read before their registering assignment (#3911)

pyrefly/lib/binding/scope.rs

@@ -1279,6 +1279,10 @@ pub struct Scope {
     has_future_annotations: bool,
     /// Tracking variables in the current scope (module, function, and method scopes)
     variables: SmallMap<Name, VariableUsage>,
+    /// Names read before any assignment registered them in `variables`.
+    /// `register_variable` treats these as already used (e.g. a loop condition
+    /// reading a name that is only assigned inside the loop body).
+    used_variables_before_registration: SmallSet<Name>,
     /// Depth of finally blocks we're in. Resets in new function scopes (PEP 765).
     finally_depth: usize,
     /// Depth of with blocks we're in. Resets in new function scopes.
@@ -1313,6 +1317,7 @@ impl Scope {
             imports: SmallMap::new(),
             has_future_annotations: false,
             variables: SmallMap::new(),
+            used_variables_before_registration: SmallSet::new(),
             finally_depth: 0,
             with_depth: 0,
             implicit_captures: SmallSet::new(),
@@ -2541,6 +2546,10 @@ impl Scopes {
                 .variables
                 .get(&name.id)
                 .is_some_and(|usage| usage.used)
+                || self
+                    .current()
+                    .used_variables_before_registration
+                    .contains(&name.id)
                 || self.is_parameter_used(&name.id);
             self.current_mut().variables.insert(
                 name.id.clone(),
@@ -2552,11 +2561,67 @@ impl Scopes {
         }
     }
 
+    /// Mark a name read as a use of the variable it resolves to, for unused-variable
+    /// tracking. The walk mirrors the scope resolution of `look_up_name_for_read`,
+    /// but is based on static scope info rather than flow: a read must count as a
+    /// use of its variable even when it is traversed before the assignment that
+    /// registers the variable (e.g. a loop condition reading a name reassigned in
+    /// the loop body). Static scope is prebuilt from the definitions pass, so it
+    /// already knows every name the scope will bind.
     pub fn mark_variable_used(&mut self, name: &Name) {
-        for scope in self.iter_rev_mut() {
+        // Class scopes are invisible to nested code blocks, except annotation
+        // scopes, which can see their enclosing class scope (see `visit_scopes`).
+        let is_current_scope_annotation_like = matches!(
+            self.current().kind,
+            ScopeKind::Annotation | ScopeKind::TypeAlias
+        );
+        for (lookup_depth, scope) in self.iter_rev_mut().enumerate() {
+            if matches!(scope.kind, ScopeKind::Class(_)) {
+                if (lookup_depth == 0 || (is_current_scope_annotation_like && lookup_depth == 1))
+                    && scope
+                        .flow
+                        .get_info(name)
+                        .is_some_and(|info| !matches!(info.initialized(), InitializedInFlow::No))
+                {
+                    // A read directly in a class body resolves to the class attribute
+                    // if one is already bound in flow (class body scopes are dynamic).
+                    // Class attributes are not tracked as variables, so there is
+                    // nothing to mark.
+                    return;
+                }
+                continue;
+            }
+
             if let Some(info) = scope.variables.get_mut(name) {
                 info.used = true;
-                break;
+                return;
+            }
+
+            if let Some(static_info) = scope.stat.0.get(name) {
+                if matches!(static_info.style, StaticStyle::MutableCapture(_)) {
+                    // `global`/`nonlocal` names resolve to (and mark) the declaring scope.
+                    continue;
+                }
+                if matches!(scope.kind, ScopeKind::Comprehension { .. })
+                    && scope.flow.get_info(name).is_none()
+                {
+                    // A walrus target is in the comprehension's static scope before the
+                    // write adds it to flow; a read before that point resolves to the
+                    // enclosing scope (mirrors `look_up_name_for_read`).
+                    continue;
+                }
+                // The name is statically local to this scope but not registered as a
+                // variable yet: remember the use so registration preserves it. Only
+                // these scope kinds track variables (see `register_variable`).
+                if matches!(
+                    scope.kind,
+                    ScopeKind::Module | ScopeKind::Function(_) | ScopeKind::Method(_)
+                ) {
+                    scope
+                        .used_variables_before_registration
+                        .insert(name.clone());
+                }
+                return;
             }
         }
     }

pyrefly/lib/test/lsp/diagnostic.rs

@@ -206,6 +206,172 @@ def test() -> Generator[float, float, None]:
     assert_eq!(report, "No unused variables");
 }
 
+#[test]
+fn test_method_read_of_name_shadowed_by_class_field() {
+    // Python name resolution skips enclosing class scopes: `return x` in the
+    // method reads f's `x`, not the class field `x`. So f's `x` is used.
+    let code = r#"
+def f():
+    x = 1
+    class C:
+        x = 2
+        def m(self) -> int:
+            return x
+    return C
+"#;
+    let (handles, state) = mk_multi_file_state(&[("main", code)], Require::Exports, true);
+    let handle = handles.get("main").unwrap();
+    let report = get_unused_variable_diagnostics(&state, handle);
+    assert_eq!(report, "No unused variables");
+}
+
+#[test]
+fn test_unpacked_target_reassigned_in_loop() {
+    // Reads of `ri` are all traversed before its only registered assignment
+    // (`li, ri = 0, 100` is an unpacking, which is not registered), so this
+    // exercises the read-before-registration tracking.
+    let code = r#"
+def foo() -> int:
+    li, ri = 0, 100
+    while li <= ri:
+        mid = (li + ri) // 2
+        if mid % 2 == 0:
+            li = mid + 1
+        else:
+            ri = mid - 1
+    return li
+"#;
+    let (handles, state) = mk_multi_file_state(&[("main", code)], Require::Exports, true);
+    let handle = handles.get("main").unwrap();
+    let report = get_unused_variable_diagnostics(&state, handle);
+    assert_eq!(report, "No unused variables");
+}
+
+#[test]
+fn test_generator_loop_reassignments() {
+    let code = r#"
+from typing import Generator
+
+def foo() -> Generator[int]:
+    li, ri = 0, 100
+    while li <= ri:
+        mid = (li + ri) // 2
+        if mid % 2 == 0:
+            li = mid + 1
+        else:
+            ri = mid - 1
+        yield li
+"#;
+    let (handles, state) = mk_multi_file_state(&[("main", code)], Require::Exports, true);
+    let handle = handles.get("main").unwrap();
+    let report = get_unused_variable_diagnostics(&state, handle);
+    assert_eq!(report, "No unused variables");
+}
+
+#[test]
+fn test_for_loop_targets_are_not_reported_as_unused_variables() {
+    let code = r#"
+def f():
+    for _ in range(3):
+        pass
+    for i in range(3):
+        x = i + 1
+        if x % 2:
+            i = 0
+        yield x
+"#;
+    let (handles, state) = mk_multi_file_state(&[("main", code)], Require::Exports, true);
+    let handle = handles.get("main").unwrap();
+    let report = get_unused_variable_diagnostics(&state, handle);
+    assert_eq!(report, "No unused variables");
+}
+
+#[test]
+fn test_with_targets_are_not_reported_as_unused_variables() {
+    let code = r#"
+def f():
+    with open("test") as handle:
+        x = handle.read()
+        handle = None
+    return x
+"#;
+    let (handles, state) = mk_multi_file_state(&[("main", code)], Require::Exports, true);
+    let handle = handles.get("main").unwrap();
+    let report = get_unused_variable_diagnostics(&state, handle);
+    assert_eq!(report, "No unused variables");
+}
+
+#[test]
+fn test_comprehension_target_does_not_mark_outer_variable_used() {
+    let code = r#"
+def f():
+    i = 0
+    result = [i for i in range(3)]
+    return result
+"#;
+    let (handles, state) = mk_multi_file_state(&[("main", code)], Require::Exports, true);
+    let handle = handles.get("main").unwrap();
+    let report = get_unused_variable_diagnostics(&state, handle);
+    assert_eq!(report, "Variable `i` is unused");
+}
+
+#[test]
+fn test_class_scope_reads_do_not_confuse_outer_variable_usage() {
+    let code = r#"
+def class_reads_initialized_class_name():
+    unused = 0
+    class C:
+        x = 1
+        y = x
+    return C
+
+def class_reads_outer_name_before_class_assignment():
+    used = 0
+    class C:
+        y = used
+        used = 1
+    return C
+"#;
+    let (handles, state) = mk_multi_file_state(&[("main", code)], Require::Exports, true);
+    let handle = handles.get("main").unwrap();
+    let report = get_unused_variable_diagnostics(&state, handle);
+    assert_eq!(report, "Variable `unused` is unused");
+}
+
+#[test]
+fn test_nonlocal_read_marks_enclosing_variable_used() {
+    let code = r#"
+def outer():
+    x = 0
+    def inner():
+        nonlocal x
+        y = x + 1
+        x = 0
+        return y
+    return inner()
+"#;
+    let (handles, state) = mk_multi_file_state(&[("main", code)], Require::Exports, true);
+    let handle = handles.get("main").unwrap();
+    let report = get_unused_variable_diagnostics(&state, handle);
+    assert_eq!(report, "No unused variables");
+}
+
+#[test]
+fn test_walrus_targets_are_not_reported_as_unused_variables() {
+    let code = r#"
+def f(xs: list[int]):
+    if (x := len(xs)) > 0:
+        y = x + 1
+        x = 0
+        return y
+    return 0
+"#;
+    let (handles, state) = mk_multi_file_state(&[("main", code)], Require::Exports, true);
+    let handle = handles.get("main").unwrap();
+    let report = get_unused_variable_diagnostics(&state, handle);
+    assert_eq!(report, "No unused variables");
+}
+
 // TODO: x = 7 should be highlighted as unused
 #[test]
 fn test_reassignment_false_negative() {