ADD COP015 check for comprehension variable one_ prefix requirement

This new check ensures that variables in comprehensions and for-loops
are prefixed with 'one_' to improve readability and consistency. It
handles list, set, dict comprehensions and generator expressions, while
properly ignoring underscore variables and partial unpacking scenarios.

Author: vrslev

src/community_of_python_flake8_plugin/checks/comprehension_one_prefix.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+import ast
+import typing
+
+from community_of_python_flake8_plugin.constants import MIN_NAME_LENGTH
+from community_of_python_flake8_plugin.violation_codes import ViolationCodes
+from community_of_python_flake8_plugin.violations import Violation
+
+
+def _is_ignored_target(target_node: ast.expr) -> bool:
+    """Check if target should be ignored (e.g., underscore variables)."""
+    # Ignore underscore variables
+    return bool(isinstance(target_node, ast.Name) and target_node.id == "_")
+
+
+@typing.final
+class COP015ComprehensionOnePrefixCheck(ast.NodeVisitor):
+    def __init__(self, syntax_tree: ast.AST) -> None:  # noqa: ARG002
+        self.violations: list[Violation] = []
+        self.syntax_tree: typing.Final[ast.AST] = syntax_tree
+
+    def visit_ListComp(self, ast_node: ast.ListComp) -> None:
+        # Validate targets in generators (the 'v' in 'for v in lst')
+        for one_comprehension in ast_node.generators:
+            if not self._is_partial_unpacking(ast_node.elt, one_comprehension.target):
+                self._validate_comprehension_target(one_comprehension.target)
+        self.generic_visit(ast_node)
+
+    def visit_SetComp(self, ast_node: ast.SetComp) -> None:
+        # Validate targets in generators (the 'v' in 'for v in lst')
+        for one_comprehension in ast_node.generators:
+            if not self._is_partial_unpacking(ast_node.elt, one_comprehension.target):
+                self._validate_comprehension_target(one_comprehension.target)
+        self.generic_visit(ast_node)
+
+    def visit_DictComp(self, ast_node: ast.DictComp) -> None:
+        # Validate targets in generators (the 'v' in 'for v in lst')
+        # For dict comprehensions, we consider both key and value
+        # key and value are both used
+        for one_comprehension in ast_node.generators:
+            if not self._is_partial_unpacking_expr_count(2, one_comprehension.target):
+                self._validate_comprehension_target(one_comprehension.target)
+        self.generic_visit(ast_node)
+
+    def visit_For(self, ast_node: ast.For) -> None:
+        # Validate target variables in regular for-loops
+        # Apply same unpacking logic as comprehensions
+        # For-loops don't have an expression that references vars
+        if not self._is_partial_unpacking_expr_count(1, ast_node.target):
+            self._validate_comprehension_target(ast_node.target)
+        self.generic_visit(ast_node)
+
+    def visit_GeneratorExp(self, ast_node: ast.GeneratorExp) -> None:
+        # Validate targets in generators (the 'v' in 'for v in lst')
+        for one_comprehension in ast_node.generators:
+            if not self._is_partial_unpacking(ast_node.elt, one_comprehension.target):
+                self._validate_comprehension_target(one_comprehension.target)
+        self.generic_visit(ast_node)
+
+    def _is_partial_unpacking(self, expression: ast.expr, target_node: ast.expr) -> bool:
+        """Check if this is partial unpacking (referencing fewer vars than unpacked)."""
+        return self._is_partial_unpacking_expr_count(self._count_referenced_vars(expression), target_node)
+
+    def _is_partial_unpacking_expr_count(self, expression_count: int, target_node: ast.expr) -> bool:
+        """Check if this is partial unpacking given expression variable count."""
+        target_count: typing.Final = self._count_unpacked_vars(target_node)
+        return target_count > expression_count and target_count > 1
+
+    def _count_referenced_vars(self, expression: ast.expr) -> int:
+        """Count how many variables are referenced in the expression."""
+        if isinstance(expression, ast.Name):
+            return 1
+        if isinstance(expression, ast.Tuple):
+            return len([one_element for one_element in expression.elts if isinstance(one_element, ast.Name)])
+        # For simplicity, assume other expressions reference 1 variable
+        # This covers cases like function calls, attributes, etc.
+        return 1
+
+    def _count_unpacked_vars(self, target_node: ast.expr) -> int:
+        """Count how many variables are unpacked in the target."""
+        if isinstance(target_node, ast.Name):
+            return 1
+        if isinstance(target_node, ast.Tuple):
+            return len([one_element for one_element in target_node.elts if isinstance(one_element, ast.Name)])
+        return 0
+
+    def _validate_comprehension_target(self, target_node: ast.expr) -> None:
+        """Validate that comprehension target follows the one_ prefix rule."""
+        # Skip ignored targets (underscore, unpacking)
+        if _is_ignored_target(target_node):
+            return
+
+        # For simple names, validate the prefix
+        if isinstance(target_node, ast.Name):
+            if not self._has_valid_one_prefix(target_node.id):
+                self.violations.append(
+                    Violation(
+                        line_number=target_node.lineno,
+                        column_number=target_node.col_offset,
+                        violation_code=ViolationCodes.COMPREHENSION_VARIABLE_PREFIX,
+                    )
+                )
+        # For tuples (unpacking), validate each element
+        elif isinstance(target_node, ast.Tuple):
+            for one_element in target_node.elts:
+                if isinstance(one_element, ast.Name) and not self._has_valid_one_prefix(one_element.id):
+                    self.violations.append(
+                        Violation(
+                            line_number=one_element.lineno,
+                            column_number=one_element.col_offset,
+                            violation_code=ViolationCodes.COMPREHENSION_VARIABLE_PREFIX,
+                        )
+                    )
+
+    def _has_valid_one_prefix(self, identifier: str) -> bool:
+        """Check if identifier has valid one_ prefix."""
+        # Allow underscore variables
+        if identifier == "_":
+            return True
+            
+        # Check for one_ prefix
+        return identifier.startswith("one_")

src/community_of_python_flake8_plugin/violation_codes.py

@@ -47,3 +47,8 @@ class ViolationCodes:
     DATACLASS_CONFIG = ViolationCodeItem(
         code="COP014", description="Use dataclasses with kw_only=True, slots=True, frozen=True"
     )
+
+    # Comprehension variable prefix violations
+    COMPREHENSION_VARIABLE_PREFIX = ViolationCodeItem(
+        code="COP015", description="Comprehension variables must be prefixed with 'one_'"
+    )

tests/test_plugin.py

@@ -483,12 +483,6 @@ def test_dataclass_validations(input_source: str, expected_output: list[str]) ->
         ("class ExampleClass:\n    a: int = 1", ["COP012", "COP004"]),
         # COP004: Class-level Final assignment should still be treated as attribute
         ("class ExampleClass:\n    client: typing.Final = CreditStatementClient()", ["COP012", "COP004"]),
-        # COP005: List comprehension with short variable names should be flagged
-        ("result_long_name = [v for v in some_list]", ["COP005"]),
-        # No violation: List comprehension with long variable names
-        ("result_long_name = [value for value in some_list]", []),
-        # No violation: List comprehension with underscore variable
-        ("result_long_name = [_ for _ in some_list]", []),
         # COP005: With statement with short variable name should be flagged
         ("with open('file.txt') as f: pass", ["COP005"]),
         # No violation: With statement with long variable name
@@ -503,6 +497,26 @@ def test_dataclass_validations(input_source: str, expected_output: list[str]) ->
         ("func_long_name = lambda value: value * 2", []),
         # COP006: Lambda with multiple short arguments should be flagged
         ("func_long_name = lambda a, b: a + b", ["COP006", "COP006"]),
+        # COP015: List comprehension without one_ prefix should be flagged (in addition to COP005)
+        ("result_long_name = [v for v in some_list]", ["COP005", "COP015"]),
+        # COP005: List comprehension with one_ prefix (still too short for general rule)
+        ("result_long_name = [one_v for one_v in some_list]", ["COP005"]),
+        # No violation: List comprehension with underscore (ignored)
+        ("result_long_name = [_ for _ in some_list]", []),
+        # COP005: List comprehension with tuple unpacking (prefix rule ignored, but length rule still applies)
+        ("result_long_name = [x for x, y in pairs]", ["COP005", "COP005"]),
+        # COP015: Set comprehension without one_ prefix should be flagged (in addition to COP005)
+        ("result_long_name = {v for v in some_list}", ["COP005", "COP015"]),
+        # COP005: Set comprehension with one_ prefix (vars still too short for general rule)
+        ("result_long_name = {one_v for one_v in some_list}", ["COP005"]),
+        # COP005+COP015: Dict comprehension without one_ prefix should be flagged
+        ("result_long_name = {k: v for k, v in pairs}", ["COP005", "COP005", "COP015", "COP015"]),
+        # COP005: Dict comprehension with one_ prefix (vars still too short for general rule)
+        ("result_long_name = {one_k: one_v for one_k, one_v in pairs}", ["COP005", "COP005"]),
+        # COP015: Generator expression without one_ prefix should be flagged (in addition to COP005)
+        ("result_long_name = (v for v in some_list)", ["COP005", "COP015"]),
+        # COP005: Generator expression with one_ prefix (vars still too short for general rule)
+        ("result_long_name = (one_v for one_v in some_list)", ["COP005"]),
     ],
 )
 def test_module_vs_class_level_assignments(input_source: str, expected_output: list[str]) -> None: