ENHANCE TEMPORARY VARIABLE CHECK WITH LINE AND ASSIGNMENT CONSTRAINTS

- Add functions to track variable assignments and check single-line
  usage
- Only flag variables assigned and used immediately in next line
- Exclude multi-line assignments and tuple unpacking from violations
- Update tests to cover new edge cases and constraints

Author: vrslev

src/community_of_python_flake8_plugin/checks/temp_var.py

@@ -19,9 +19,28 @@ def extract_names(expression: ast.expr) -> typing.Iterable[str]:
             yield from extract_names(elt)
 
 
-def collect_variable_usage_and_stores(function_node: ast.AST) -> tuple[dict[str, list[ast.Name]], set[str]]:
+def is_single_line_assignment(assign_node: ast.Assign) -> bool:
+    """Check if assignment value fits on a single line."""
+    # For simple values, they're always single line
+    if isinstance(assign_node.value, (ast.Constant, ast.Name, ast.Attribute)):
+        return True
+    
+    # For complex expressions, check if they span multiple lines
+    # We'll consider it multi-line if the end line is different from start line
+    if hasattr(assign_node, 'lineno') and hasattr(assign_node.value, 'end_lineno'):
+        return assign_node.lineno == assign_node.value.end_lineno
+    
+    return True
+
+
+def collect_variable_usage_and_stores_with_nodes(function_node: ast.AST) -> tuple[
+    dict[str, list[ast.Name]], 
+    set[str], 
+    dict[str, ast.Assign]
+]:
     variable_usage: typing.Final[defaultdict[str, list[ast.Name]]] = defaultdict(list)
     assigned_variable_names: typing.Final[set[str]] = set()
+    variable_assignments: typing.Final[dict[str, ast.Assign]] = {}
 
     @typing.final
     class UsageCollector(ast.NodeVisitor):
@@ -36,19 +55,44 @@ def visit_Assign(self, assign_node: ast.Assign) -> None:
                 return
 
             for target in assign_node.targets:
-                assigned_variable_names.update(extract_names(target))
+                names = list(extract_names(target))
+                assigned_variable_names.update(names)
+                # Store the assignment node for each variable
+                for name in names:
+                    variable_assignments[name] = assign_node
             self.generic_visit(assign_node)
 
         def visit_AugAssign(self, aug_assign_node: ast.AugAssign) -> None:
-            assigned_variable_names.update(extract_names(aug_assign_node.target))
+            names = list(extract_names(aug_assign_node.target))
+            assigned_variable_names.update(names)
             self.generic_visit(aug_assign_node)
 
         def visit_AnnAssign(self, ann_assign_node: ast.AnnAssign) -> None:
-            assigned_variable_names.update(extract_names(ann_assign_node.target))
+            names = list(extract_names(ann_assign_node.target))
+            assigned_variable_names.update(names)
+            # Store the assignment node for each variable
+            for name in names:
+                variable_assignments[name] = ann_assign_node
             self.generic_visit(ann_assign_node)
 
     UsageCollector().visit(function_node)
-    return dict(variable_usage), assigned_variable_names
+    return dict(variable_usage), assigned_variable_names, variable_assignments
+
+
+def is_used_in_next_line(assign_node: ast.Assign, usage_nodes: list[ast.Name]) -> bool:
+    """Check if variable is used in the immediate next line."""
+    if not usage_nodes:
+        return False
+        
+    # Find the load usage (actual use of the variable)
+    load_usages = [node for node in usage_nodes if isinstance(node.ctx, ast.Load)]
+    if not load_usages:
+        return False
+        
+    first_load = load_usages[0]
+    
+    # Check if the usage is on the line immediately following the assignment
+    return first_load.lineno == assign_node.lineno + 1
 
 
 @typing.final
@@ -65,25 +109,40 @@ def visit_AsyncFunctionDef(self, ast_node: ast.AsyncFunctionDef) -> None:
         self.generic_visit(ast_node)
 
     def _check_temporary_variables(self, ast_node: ast.FunctionDef | ast.AsyncFunctionDef) -> None:
-        usage_and_stores: typing.Final = collect_variable_usage_and_stores(ast_node)
-        found_temporary_variable = False
+        usage_and_stores: typing.Final = collect_variable_usage_and_stores_with_nodes(ast_node)
+        variable_usages = usage_and_stores[0]
+        assigned_variable_names = usage_and_stores[1]
+        variable_assignments = usage_and_stores[2]
 
-        for variable_name, usages in usage_and_stores[0].items():
+        for variable_name, usages in variable_usages.items():
             if variable_name.startswith("_") or variable_name in {"self", "cls"}:
                 continue
 
-            if (
-                len([usage_element for usage_element in usages if isinstance(usage_element.ctx, ast.Store)]) == 1
-                and len([usage_element for usage_element in usages if isinstance(usage_element.ctx, ast.Load)]) == 1
-                and variable_name in usage_and_stores[1]
-                and not found_temporary_variable
-            ):
-                first_usage = usages[0]
-                self.violations.append(
-                    Violation(
-                        line_number=first_usage.lineno,
-                        column_number=first_usage.col_offset,
-                        violation_code=ViolationCodes.TEMP_VAR,
-                    )
-                )
-                found_temporary_variable = True
+            store_count = len([usage_element for usage_element in usages if isinstance(usage_element.ctx, ast.Store)])
+            load_count = len([usage_element for usage_element in usages if isinstance(usage_element.ctx, ast.Load)])
+
+            # Check if variable is assigned once and used once
+            if store_count == 1 and load_count == 1 and variable_name in assigned_variable_names:
+                # Get the assignment node
+                assign_node = variable_assignments.get(variable_name)
+                if not assign_node:
+                    continue
+                    
+                # Check if it's a single-line assignment
+                if not is_single_line_assignment(assign_node):
+                    continue
+                    
+                # Check if it's used in the next line
+                if is_used_in_next_line(assign_node, usages):
+                    # Only flag the variable that is assigned and then immediately used
+                    # Find the store usage (assignment)
+                    store_usages = [node for node in usages if isinstance(node.ctx, ast.Store)]
+                    if store_usages:
+                        first_store = store_usages[0]
+                        self.violations.append(
+                            Violation(
+                                line_number=first_store.lineno,
+                                column_number=first_store.col_offset,
+                                violation_code=ViolationCodes.TEMP_VAR,
+                            )
+                        )

tests/test_cop011_modified.py

@@ -0,0 +1,69 @@
+import ast
+import textwrap
+
+from community_of_python_flake8_plugin.checks.temp_var import TempVarCheck
+
+
+def test_cop011_applies_to_single_line_immediate_usage():
+    """Test that COP011 applies to single-line variables used in next line."""
+    code = textwrap.dedent("""
+        def func():
+            a = 1
+            b = a + 1
+    """)
+    
+    tree = ast.parse(code)
+    checker = TempVarCheck(tree)
+    checker.visit(tree)
+    
+    assert len(checker.violations) == 1
+    assert checker.violations[0].violation_code.code == "COP011"
+
+
+def test_cop011_not_applied_with_intervening_lines():
+    """Test that COP011 does not apply when there are intervening lines."""
+    code = textwrap.dedent("""
+        def func():
+            a = 1
+            # some things here
+            b = a + 1
+    """)
+    
+    tree = ast.parse(code)
+    checker = TempVarCheck(tree)
+    checker.visit(tree)
+    
+    assert len(checker.violations) == 0
+
+
+def test_cop011_not_applied_to_multi_line_assignment():
+    """Test that COP011 does not apply to multi-line assignments."""
+    code = textwrap.dedent("""
+        def func():
+            a = (
+                one_item for one_item
+                in lst
+            )
+            b = a[0]
+    """)
+    
+    tree = ast.parse(code)
+    checker = TempVarCheck(tree)
+    checker.visit(tree)
+    
+    assert len(checker.violations) == 0
+
+
+def test_cop011_not_applied_to_tuple_unpacking():
+    """Test that COP011 does not apply to tuple unpacking."""
+    code = textwrap.dedent("""
+        def func():
+            a, b = (1, 2)
+            c = a + b
+    """)
+    
+    tree = ast.parse(code)
+    checker = TempVarCheck(tree)
+    checker.visit(tree)
+    
+    assert len(checker.violations) == 0

tests/test_plugin.py

@@ -280,6 +280,11 @@ def test_naming_validations(input_source: str, expected_output: list[str]) -> No
         (
             "def fetch_value() -> int:\n    result_value = 1\n    another_value = result_value\n    "
             "return another_value",
+            ["COP011", "COP011"],
+        ),
+        # COP011: Single-line assignment used immediately in next line
+        (
+            "def func():\n    a = 1\n    b = a + 1",
             ["COP011"],
         ),
         # No violation: Variable used multiple times
@@ -296,6 +301,21 @@ def test_naming_validations(input_source: str, expected_output: list[str]) -> No
         ("def fetch_item():\n    for _, one_value in values: print(f'{one_value}')", []),
         # No violation: Tuple unpacking with underscore variables
         ("def parse_module():\n    parent, _, _child = module_name.rpartition('.')\n    return parent", []),
+        # No violation: Single-line assignment with intervening lines
+        (
+            "def func():\n    a = 1\n    # some things here\n    b = a + 1",
+            [],
+        ),
+        # No violation: Multi-line assignment used immediately
+        (
+            "def func():\n    a = (\n        one_item for one_item\n        in lst\n    )\n    b = a[0]",
+            [],
+        ),
+        # No violation: Tuple unpacking assignment
+        (
+            "def func():\n    a, b = (1, 2)\n    c = a + b",
+            [],
+        ),
     ],
 )
 def test_variable_usage_validations(input_source: str, expected_output: list[str]) -> None: