Implement ForVariableUsedAfterBlock (based on wemake_python_styleguide ControlVarUsedAfterBlockViolation WPS441)

crates/ruff_linter/resources/test/fixtures/ruff/for_variable_used_after_block.py

@@ -0,0 +1,73 @@
+for global_var in []:
+    _ = global_var
+    pass
+
+# ❌
+_ = global_var
+
+def foo():
+    # For control var used outside block
+    for event in []:
+        _ = event
+        pass
+
+    # ❌
+    _ = event
+
+    for x in []:
+        _ = x
+        pass
+
+    # Using the same control variable name in a different loop is ok
+    for x in []:
+        _ = x
+        pass
+
+    for y in []:
+        # ❌ x is used outside of the loop it was defined in (meant to use y)
+        if x == 5:
+            pass
+
+    # Assign a variable before the loop
+    room_id = 3
+    _ = room_id
+    # Use the same variable name in a loop
+    for room_id in []:
+        _ = room_id
+        pass
+
+    # ❌ After the loop is not ok because the value is probably not what you expect
+    _ = room_id
+
+    # Tuple destructuring
+    for a, b, c in []:
+        _ = a
+        _ = b
+        _ = c
+        pass
+
+    # ❌
+    _ = a
+    _ = b
+    _ = c
+
+    # Array destructuring
+    for [d, e, f] in []:
+        _ = d
+        _ = e
+        _ = f
+        pass
+
+    # ❌
+    _ = d
+    _ = e
+    _ = f
+
+    # Nested function and class definitions are fine
+    for potential_power in []:
+        def action():
+            print(potential_power)
+
+        class Animal:
+            power_level = potential_power
+

crates/ruff_linter/src/checkers/ast/analyze/deferred_scopes.rs

@@ -17,6 +17,7 @@ pub(crate) fn deferred_scopes(checker: &mut Checker) {
         Rule::AsyncioDanglingTask,
         Rule::BadStaticmethodArgument,
         Rule::BuiltinAttributeShadowing,
+        Rule::ForVariableUsedAfterBlock,
         Rule::GlobalVariableNotAssigned,
         Rule::ImportPrivateName,
         Rule::ImportShadowedByLoopVar,
@@ -86,6 +87,10 @@ pub(crate) fn deferred_scopes(checker: &mut Checker) {
     for scope_id in checker.analyze.scopes.iter().rev().copied() {
         let scope = &checker.semantic.scopes[scope_id];
 
+        if checker.enabled(Rule::ForVariableUsedAfterBlock) {
+            ruff::rules::for_variable_used_after_block(checker, scope, &mut diagnostics);
+        }
+
         if checker.enabled(Rule::UndefinedLocal) {
             pyflakes::rules::undefined_local(checker, scope_id, scope, &mut diagnostics);
         }

crates/ruff_linter/src/codes.rs

@@ -955,6 +955,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Ruff, "028") => (RuleGroup::Preview, rules::ruff::rules::InvalidFormatterSuppressionComment),
         (Ruff, "029") => (RuleGroup::Preview, rules::ruff::rules::UnusedAsync),
         (Ruff, "030") => (RuleGroup::Preview, rules::ruff::rules::AssertWithPrintMessage),
+        (Ruff, "031") => (RuleGroup::Preview, rules::ruff::rules::ForVariableUsedAfterBlock),
         (Ruff, "100") => (RuleGroup::Stable, rules::ruff::rules::UnusedNOQA),
         (Ruff, "101") => (RuleGroup::Preview, rules::ruff::rules::RedirectedNOQA),
 

crates/ruff_linter/src/rules/ruff/mod.rs

@@ -24,6 +24,10 @@ mod tests {
 
     #[test_case(Rule::AsyncioDanglingTask, Path::new("RUF006.py"))]
     #[test_case(Rule::CollectionLiteralConcatenation, Path::new("RUF005.py"))]
+    #[test_case(
+        Rule::ForVariableUsedAfterBlock,
+        Path::new("for_variable_used_after_block.py")
+    )]
     #[test_case(Rule::ExplicitFStringTypeConversion, Path::new("RUF010.py"))]
     #[test_case(Rule::FunctionCallInDataclassDefaultArgument, Path::new("RUF009.py"))]
     #[test_case(Rule::ImplicitOptional, Path::new("RUF013_0.py"))]

crates/ruff_linter/src/rules/ruff/rules/for_variable_used_after_block.rs

@@ -0,0 +1,109 @@
+use ruff_diagnostics::{Diagnostic, Violation};
+use ruff_macros::{derive_message_formats, violation};
+use ruff_python_semantic::{BindingId, NodeId, Scope};
+use ruff_text_size::Ranged;
+
+use crate::checkers::ast::Checker;
+
+/// ## What it does
+/// Checks for variables defined in `for` loops that are used outside of their
+/// respective blocks.
+///
+/// ## Why is this bad?
+/// Usage of of a control variable outside of the block they're defined in will probably
+/// lead to flawed logic in a way that will likely cause bugs. The variable might not
+/// contain what you expect.
+///
+/// In Python, unlike many other languages, `for` loops don't define their own scopes.
+/// Therefore, usage of the control variables outside of the block will be the the value
+/// from the last iteration until re-assigned.
+///
+/// While this mistake is easy to spot in small examples, it can be hidden in larger
+/// blocks of code, where the the loop and downstream usage may not be visible at the
+/// same time.
+///
+/// ## Example
+/// ```python
+/// for x in range(10):
+///     pass
+///
+/// print(x)  # prints 9
+/// ```
+#[violation]
+pub struct ForVariableUsedAfterBlock {
+    control_var_name: String,
+}
+
+impl Violation for ForVariableUsedAfterBlock {
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        let ForVariableUsedAfterBlock { control_var_name } = self;
+
+        format!("`for` loop variable {control_var_name} is used outside of block")
+    }
+}
+
+/// Based on wemake-python-styleguide (WPS441) to forbid control variables after the block body.
+pub(crate) fn for_variable_used_after_block(
+    checker: &Checker,
+    scope: &Scope,
+    diagnostics: &mut Vec<Diagnostic>,
+) {
+    // Keep track of the node_ids of variable references that have already been
+    // accounted for as part of a different variable binding. This helps us avoid
+    // complaining when people use the same variable name in multiple blocks under the
+    // same scope.
+    let mut known_good_reference_node_ids: Vec<NodeId> = Vec::new();
+
+    let all_bindings: Vec<(&str, BindingId)> = scope.all_bindings().collect();
+    // We want to reverse the bindings so that we iterate in source order and shadowed
+    // bindings come first. This way we can gather `known_good_reference_node_ids` and
+    // mark off things as we go which will allow us to ignore later bindings trying to
+    // reference the same variable.
+    let reversed_bindings = all_bindings.iter().rev();
+
+    // Find for-loop variable bindings
+    let loop_var_bindings = reversed_bindings
+        .map(|(name, binding_id)| (name, checker.semantic().binding(*binding_id)))
+        .filter_map(|(name, binding)| binding.kind.is_loop_var().then_some((name, binding)));
+
+    for (&name, binding) in loop_var_bindings {
+        let binding_statement = binding.statement(checker.semantic()).unwrap();
+        let binding_source_node_id = binding.source.unwrap();
+        // The node_id of the for-loop that contains the binding
+        let binding_statement_id = checker.semantic().statement_id(binding_source_node_id);
+
+        // Loop over the references of those bindings to see if they're in the same block-scope
+        'references: for reference in binding.references() {
+            let reference = checker.semantic().reference(reference);
+            let reference_node_id = reference.expression_id().unwrap();
+
+            // Skip any reference that come before the control var binding in the source
+            // order, skip it because people can assign and use the same variable name
+            // above the block.
+            if reference.range().end() < binding_statement.range().start() {
+                continue;
+            }
+
+            // Traverse the hierarchy and look for a block match
+            let statement_hierarchy = checker.semantic().parent_statement_ids(reference_node_id);
+
+            for ancestor_node_id in statement_hierarchy {
+                if binding_statement_id == ancestor_node_id {
+                    known_good_reference_node_ids.push(reference_node_id);
+                    continue 'references;
+                }
+            }
+
+            // If the reference wasn't used in the same block, report a violation/diagnostic
+            if !known_good_reference_node_ids.contains(&reference_node_id) {
+                diagnostics.push(Diagnostic::new(
+                    ForVariableUsedAfterBlock {
+                        control_var_name: name.to_owned(),
+                    },
+                    reference.range(),
+                ));
+            }
+        }
+    }
+}

crates/ruff_linter/src/rules/ruff/rules/mod.rs

@@ -5,6 +5,7 @@ pub(crate) use asyncio_dangling_task::*;
 pub(crate) use collection_literal_concatenation::*;
 pub(crate) use default_factory_kwarg::*;
 pub(crate) use explicit_f_string_type_conversion::*;
+pub(crate) use for_variable_used_after_block::*;
 pub(crate) use function_call_in_dataclass_default::*;
 pub(crate) use implicit_optional::*;
 pub(crate) use invalid_formatter_suppression_comment::*;
@@ -38,6 +39,7 @@ mod collection_literal_concatenation;
 mod confusables;
 mod default_factory_kwarg;
 mod explicit_f_string_type_conversion;
+mod for_variable_used_after_block;
 mod function_call_in_dataclass_default;
 mod helpers;
 mod implicit_optional;

crates/ruff_linter/src/rules/ruff/snapshots/ruff_linter__rules__ruff__tests__RUF031_for_variable_used_after_block.py.snap

@@ -0,0 +1,94 @@
+---
+source: crates/ruff_linter/src/rules/ruff/mod.rs
+---
+for_variable_used_after_block.py:6:5: RUF031 `for` loop variable global_var is used outside of block
+  |
+5 | # ❌
+6 | _ = global_var
+  |     ^^^^^^^^^^ RUF031
+7 | 
+8 | def foo():
+  |
+
+for_variable_used_after_block.py:15:9: RUF031 `for` loop variable event is used outside of block
+   |
+14 |     # ❌
+15 |     _ = event
+   |         ^^^^^ RUF031
+16 | 
+17 |     for x in []:
+   |
+
+for_variable_used_after_block.py:28:12: RUF031 `for` loop variable x is used outside of block
+   |
+26 |     for y in []:
+27 |         # ❌ x is used outside of the loop it was defined in (meant to use y)
+28 |         if x == 5:
+   |            ^ RUF031
+29 |             pass
+   |
+
+for_variable_used_after_block.py:40:9: RUF031 `for` loop variable room_id is used outside of block
+   |
+39 |     # ❌ After the loop is not ok because the value is probably not what you expect
+40 |     _ = room_id
+   |         ^^^^^^^ RUF031
+41 | 
+42 |     # Tuple destructuring
+   |
+
+for_variable_used_after_block.py:50:9: RUF031 `for` loop variable a is used outside of block
+   |
+49 |     # ❌
+50 |     _ = a
+   |         ^ RUF031
+51 |     _ = b
+52 |     _ = c
+   |
+
+for_variable_used_after_block.py:51:9: RUF031 `for` loop variable b is used outside of block
+   |
+49 |     # ❌
+50 |     _ = a
+51 |     _ = b
+   |         ^ RUF031
+52 |     _ = c
+   |
+
+for_variable_used_after_block.py:52:9: RUF031 `for` loop variable c is used outside of block
+   |
+50 |     _ = a
+51 |     _ = b
+52 |     _ = c
+   |         ^ RUF031
+53 | 
+54 |     # Array destructuring
+   |
+
+for_variable_used_after_block.py:62:9: RUF031 `for` loop variable d is used outside of block
+   |
+61 |     # ❌
+62 |     _ = d
+   |         ^ RUF031
+63 |     _ = e
+64 |     _ = f
+   |
+
+for_variable_used_after_block.py:63:9: RUF031 `for` loop variable e is used outside of block
+   |
+61 |     # ❌
+62 |     _ = d
+63 |     _ = e
+   |         ^ RUF031
+64 |     _ = f
+   |
+
+for_variable_used_after_block.py:64:9: RUF031 `for` loop variable f is used outside of block
+   |
+62 |     _ = d
+63 |     _ = e
+64 |     _ = f
+   |         ^ RUF031
+65 | 
+66 |     # Nested function and class definitions are fine
+   |

crates/ruff_python_semantic/src/model.rs

@@ -1201,9 +1201,17 @@ impl<'a> SemanticModel<'a> {
     /// Return the [`Stmt`] corresponding to the given [`NodeId`].
     #[inline]
     pub fn statement(&self, node_id: NodeId) -> &'a Stmt {
+        self.nodes[self.statement_id(node_id)]
+            .as_statement()
+            .expect("No statement found")
+    }
+
+    /// Return the statement [`NodeId`] corresponding to the given [`NodeId`].
+    #[inline]
+    pub fn statement_id(&self, node_id: NodeId) -> NodeId {
         self.nodes
             .ancestor_ids(node_id)
-            .find_map(|id| self.nodes[id].as_statement())
+            .find(|id| self.nodes[*id].is_statement())
             .expect("No statement found")
     }
 
@@ -1232,6 +1240,13 @@ impl<'a> SemanticModel<'a> {
             .nth(1)
     }
 
+    /// Given a [`NodeId`], return the [`NodeId`] of the parent statement, if any.
+    pub fn parent_statement_ids(&self, node_id: NodeId) -> impl Iterator<Item = NodeId> + '_ {
+        self.nodes
+            .ancestor_ids(node_id)
+            .filter(|id| self.nodes[*id].is_statement())
+    }
+
     /// Return the [`Expr`] corresponding to the given [`NodeId`].
     #[inline]
     pub fn expression(&self, node_id: NodeId) -> Option<&'a Expr> {