fix: leaking for loops, refs wemake-services#2340 (wemake-services#3532)

Co-authored-by: Alexey Potapov <[email protected] >

Author: milssky

CHANGELOG.md

@@ -17,6 +17,13 @@ Semantic versioning in our case means:
   change the client facing API, change code conventions significantly, etc.
 
 
+## WIP
+
+### Features
+
+- Adds `WPS481`: for statement not allowed in class and module scopes, #3520
+
+
 ## 1.5.0
 
 ### Features

tests/fixtures/noqa/noqa.py

@@ -317,7 +317,7 @@ def function_with_wrong_yield():
         yield  # noqa: WPS325
     yield 1
 
-for literal in bad_concatenation:  # noqa: WPS327, WPS328
+for literal in bad_concatenation:  # noqa: WPS327, WPS328, WPS481
     continue
 
 with open(bad_concatenation):  # noqa: WPS328
@@ -326,7 +326,7 @@ def function_with_wrong_yield():
 
 my_print(biggesst > middle >= smallest)  # noqa: WPS334
 
-for index in [1, 2]:  # noqa: WPS335
+for index in [1, 2]:  # noqa: WPS335, WPS481
     my_print(index)
 
 string_concat = 'a' + 'b'  # noqa: WPS336
@@ -339,11 +339,11 @@ def function_with_wrong_yield():
 my_print(lambda: 0)  # noqa: WPS522
 xterm += xterm + 1  # noqa: WPS524
 
-for range_len in range(len(file_obj)):  # noqa: WPS518
+for range_len in range(len(file_obj)):  # noqa: WPS518, WPS481
     my_print(range_len)
 
 sum_container = 0
-for sum_item in file_obj:  # noqa: WPS519
+for sum_item in file_obj:  # noqa: WPS519, WPS481
     sum_container += sum_item
 
 my_print(sum_container == [])  # noqa: WPS520
@@ -364,7 +364,7 @@ def __init__(self) -> None:
         self.second = 2  # noqa: WPS601
 
 
-for symbol in 'abc':  # noqa: WPS500
+for symbol in 'abc':  # noqa: WPS500, WPS481
     ...
 else:
     ...
@@ -386,7 +386,7 @@ def __eq__(self, object_: object) -> bool:  # noqa: WPS612
         return super().__eq__(object_)
 
 
-for loop_index in range(6):  # noqa: WPS426
+for loop_index in range(6):  # noqa: WPS426, WPS481
     my_print(lambda: loop_index)
 
 
@@ -409,7 +409,7 @@ class MyBadException(BaseException):  # noqa: WPS418
 class ClassWithWrongContents((lambda: object)()):  # noqa: WPS606
     __slots__ = ['a', 'a']  # noqa: WPS607
 
-    for bad_body_node in range(1):  # noqa: WPS604
+    for bad_body_node in range(1):  # noqa: WPS481, WPS604
         anti_wps604 = 1
 
     def method_with_no_args():  # noqa: WPS605
@@ -424,7 +424,7 @@ def bad_default_values(
     return True
 
 
-for nodes[0] in (1, 2, 3):  # noqa: WPS405
+for nodes[0] in (1, 2, 3):  # noqa: WPS405, WPS481
     ...
 
 with open('some') as MyBadException.custom:  # noqa: WPS406
@@ -631,7 +631,7 @@ def many_raises_function(parameter):  # noqa: WPS238
 def get_item():  # noqa: WPS463
     return  # noqa: WPS324
 
-for _, something in enumerate(collection): # noqa: WPS468
+for _, something in enumerate(collection): # noqa: WPS468, WPS481
     report(something)
 
 variable_to_store_things = {
@@ -729,7 +729,7 @@ def wrong_comprehension2():
     ]
 
 
-for unique_element in range(10):
+for unique_element in range(10):  # noqa: WPS481
     if (other := unique_element) > 5:  # noqa: WPS332
         my_print(1)
 

tests/test_checker/test_noqa.py

@@ -248,6 +248,7 @@
         'WPS478': 1,
         'WPS479': 0,
         'WPS480': 0,  # only triggers on 3.12+
+        'WPS481': 10,
         'WPS500': 1,
         'WPS501': 1,
         'WPS502': 0,  # disabled since 1.0.0

tests/test_visitors/test_ast/test_conditions/test_leaking_for_loops.py

@@ -0,0 +1,175 @@
+import pytest
+
+from wemake_python_styleguide.violations.best_practices import (
+    LeakingForLoopViolation,
+)
+from wemake_python_styleguide.visitors.ast.conditions import (
+    LeakingForLoopVisitor,
+)
+
+# Wrong:
+
+module_scope_code = """
+for index in range(10):
+    ...
+"""
+
+class_body_code = """
+class ClassWithBody:
+    for index in range(10):
+        ...
+"""
+
+for_with_wrong_del = """
+for index in range(10):
+    ...
+del other
+"""
+
+for_unpacking_with_partial_del = """
+for a, b in [(1, 2)]:
+    ...
+del a
+"""
+
+separate_classes_false_cleanup = """
+class ClassWithFor:
+    for index in range(10):
+        ...
+
+class ClassWithDel:
+    index = ...
+    del index
+"""
+
+multiple_for_mixed = """
+for good in range(5):
+    ...
+del good
+
+for bad in range(5):
+    ...
+"""
+
+for_unpacking_with_unrelated_del = """
+for a, b in [(1, 2)]:
+    ...
+
+def func(a, b):
+    del a, b
+"""
+
+
+class_for_with_del_inside_method = """
+class ClassWithFor:
+    for a in range(10):
+        ...
+
+    def cleanup(self):
+        del a
+"""
+
+
+# Correct
+
+for_with_del_module = """
+for index in range(10):
+    ...
+del index
+"""
+
+for_with_del_class = """
+class ClassWithDel:
+    for index in range(10):
+        ...
+    del index
+"""
+
+for_unpacking_with_del = """
+for a, b in [(1, 2)]:
+    ...
+del a, b
+"""
+
+for_with_del_inside_if = """
+for i in range(10):
+    ...
+
+if True:
+    del i
+"""
+
+for_unpacking_with_nested_del = """
+for a, b in [(1, 2)]:
+    ...
+
+try:
+    del a, b
+except ...:
+    ...
+"""
+
+class_unpacking_with_nested_del = """
+class ClassWithNestedDel:
+    for a, b in [(1, 2)]:
+        ...
+
+    try:
+        del a, b
+    except ...:
+        ...
+"""
+
+
+@pytest.mark.parametrize(
+    'code',
+    [
+        module_scope_code,
+        class_body_code,
+        for_with_wrong_del,
+        for_unpacking_with_partial_del,
+        separate_classes_false_cleanup,
+        multiple_for_mixed,
+        for_unpacking_with_unrelated_del,
+        class_for_with_del_inside_method,
+    ],
+)
+def test_leaking_for_loop_violation(
+    assert_errors,
+    parse_ast_tree,
+    default_options,
+    code,
+):
+    """Ensure leaking for loops without proper cleanup raise a violation."""
+    tree = parse_ast_tree(code)
+
+    visitor = LeakingForLoopVisitor(default_options, tree)
+    visitor.run()
+
+    assert_errors(visitor, [LeakingForLoopViolation])
+
+
+@pytest.mark.parametrize(
+    'code',
+    [
+        for_with_del_module,
+        for_with_del_class,
+        for_unpacking_with_del,
+        for_with_del_inside_if,
+        for_unpacking_with_nested_del,
+        class_unpacking_with_nested_del,
+    ],
+)
+def test_for_loop_with_del_no_violation(
+    assert_errors,
+    parse_ast_tree,
+    default_options,
+    code,
+):
+    """Ensure `for` loops with proper `del` do not raise a violation."""
+    tree = parse_ast_tree(code)
+
+    visitor = LeakingForLoopVisitor(default_options, tree)
+    visitor.run()
+
+    assert_errors(visitor, [])

wemake_python_styleguide/logic/walk/__init__.py

@@ -0,0 +1,24 @@
+from wemake_python_styleguide.logic.walk.deletion import (
+    are_variables_deleted as are_variables_deleted,
+)
+from wemake_python_styleguide.logic.walk.deletion import (
+    extract_deleted_names as extract_deleted_names,
+)
+from wemake_python_styleguide.logic.walk.targets import (
+    extract_names_from_targets as extract_names_from_targets,
+)
+from wemake_python_styleguide.logic.walk.targets import (
+    get_names_from_target as get_names_from_target,
+)
+from wemake_python_styleguide.logic.walk.tree import (
+    get_closest_parent as get_closest_parent,
+)
+from wemake_python_styleguide.logic.walk.tree import (
+    get_subnodes_by_type as get_subnodes_by_type,
+)
+from wemake_python_styleguide.logic.walk.tree import (
+    is_contained as is_contained,
+)
+from wemake_python_styleguide.logic.walk.tree import (
+    is_contained_by as is_contained_by,
+)

wemake_python_styleguide/logic/walk/deletion.py

@@ -0,0 +1,37 @@
+import ast
+
+from wemake_python_styleguide.logic.nodes import get_context
+from wemake_python_styleguide.logic.walk.targets import (
+    extract_names_from_targets,
+)
+from wemake_python_styleguide.types import ContextNodes
+
+
+def extract_deleted_names(
+    node: ast.AST, *, context: ContextNodes | None = None
+) -> set[str]:
+    """Extract all variable names deleted in the given AST node."""
+    deleted: set[str] = set()
+
+    for subnode in ast.walk(node):
+        if not isinstance(subnode, ast.Delete):
+            continue
+
+        if context is not None and get_context(subnode) != context:
+            continue
+
+        deleted.update(extract_names_from_targets(subnode.targets))
+    return deleted
+
+
+def are_variables_deleted(
+    variables: set[str],
+    body: list[ast.stmt],
+    *,
+    context: ContextNodes | None = None,
+) -> bool:
+    """Checks that given variables are deleted somewhere in the body."""
+    deleted: set[str] = set()
+    for stmt in body:
+        deleted.update(extract_deleted_names(stmt, context=context))
+    return variables.issubset(deleted)

wemake_python_styleguide/logic/walk/targets.py

@@ -0,0 +1,22 @@
+import ast
+
+
+def extract_names_from_targets(targets: list[ast.expr]) -> set[str]:
+    """Extract names from delete statement targets."""
+    names: set[str] = set()
+    for target in targets:
+        if isinstance(target, ast.Name):
+            names.add(target.id)
+    return names
+
+
+def get_names_from_target(target: ast.expr) -> set[str]:
+    """
+    Extracts all variable names from a target expression.
+
+    Works with simple names and unpacking, e.g.:
+    - `for x in ...` -> {"x"}
+    - `for x, y in ...` -> {"x", "y"}
+    - `for (a, (b, c)) in ...` -> {"a", "b", "c"}
+    """
+    return {node.id for node in ast.walk(target) if isinstance(node, ast.Name)}

wemake_python_styleguide/logic/walk.py renamed to wemake_python_styleguide/logic/walk/tree.py

@@ -71,6 +71,7 @@
     conditions.MatchVisitor,
     conditions.ChainedIsVisitor,
     conditions.SimplifiableMatchVisitor,
+    conditions.LeakingForLoopVisitor,
     iterables.IterableUnpackingVisitor,
     blocks.AfterBlockVariablesVisitor,
     subscripts.SubscriptVisitor,