add rule prefer-simpler-iterator

flake8_pie/__init__.py

@@ -36,6 +36,10 @@
     pie803_prefer_logging_interpolation,
 )
 from flake8_pie.pie804_no_unnecessary_dict_kwargs import pie804_no_dict_kwargs
+from flake8_pie.pie805_prefer_simple_iterator import (
+    pie805_prefer_simple_iterator_for,
+    pie805_prefer_simple_iterator_generator,
+)
 
 
 @dataclass(frozen=True)
@@ -61,6 +65,7 @@ def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
     def visit_For(self, node: ast.For) -> None:
         self._visit_body(node)
         self._visit_body(BodyNode(node.orelse))
+        pie805_prefer_simple_iterator_for(node, self.errors)
         self.generic_visit(node)
 
     def visit_AsyncFor(self, node: ast.AsyncFor) -> None:
@@ -113,6 +118,11 @@ def visit_Expr(self, node: ast.Expr) -> None:
 
         self.generic_visit(node)
 
+    def visit_GeneratorExp(self, node: ast.GeneratorExp) -> None:
+        pie805_prefer_simple_iterator_generator(node, self.errors)
+
+        self.generic_visit(node)
+
     def visit_If(self, node: ast.If) -> None:
         pie787_no_len_condition(node, self.errors)
         pie789_prefer_isinstance_type_compare(node, self.errors)
@@ -147,6 +157,11 @@ def _visit_body(self, node: Body) -> None:
         pie799_prefer_col_init(node, self.errors)
         pie801_prefer_simple_return(node, self.errors)
 
+    def visit_ListComp(self, node: ast.ListComp) -> None:
+        pie805_prefer_simple_iterator_generator(node, self.errors)
+
+        self.generic_visit(node)
+
     def visit_Module(self, node: ast.Module) -> None:
         self._visit_body(node)
         self.generic_visit(node)

flake8_pie/pie805_prefer_simple_iterator.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import ast
+import string
+from collections.abc import Iterable
+from typing import NamedTuple
+
+from typing_extensions import TypeGuard
+
+from flake8_pie.base import Error
+
+KNOWN_FUNCTIONS = {"items"}
+
+
+def is_name_list(val: Iterable[ast.expr]) -> TypeGuard[Iterable[ast.Name]]:
+    return all(isinstance(x, ast.Name) for x in val)
+
+
+class UsedVarRes(NamedTuple):
+    method: str
+    var: str
+
+
+def get_used_var(vars: Iterable[ast.Name]) -> UsedVarRes | None:
+    for idx, var in enumerate(vars):
+        if not var.id.startswith("_"):
+            if idx == 0:
+                method = "keys"
+            elif idx == 1:
+                method = "values"
+            else:
+                raise ValueError(f"Unexpected index: {idx}")
+
+            return UsedVarRes(method=method, var=var.id)
+    return None
+
+
+def has_unused_var(vars: Iterable[ast.Name]) -> bool:
+    return any(x.id.startswith("_") for x in vars)
+
+
+def pie805_prefer_simple_iterator_for(node: ast.For, errors: list[Error]) -> None:
+    if (
+        isinstance(node.target, ast.Tuple)
+        and len(node.target.elts) == 2
+        and is_name_list(node.target.elts)
+        and has_unused_var(node.target.elts)
+    ):
+        res = get_used_var(node.target.elts)
+        if res is None:
+            return
+
+        if (
+            isinstance(node.iter, ast.Call)
+            and isinstance(node.iter.func, ast.Attribute)
+            and node.iter.func.attr == "items"
+        ):
+            errors.append(
+                PIE805(
+                    lineno=node.lineno,
+                    col_offset=node.col_offset,
+                    suggestion=f"use `for {res.var} in foo.{res.method}()`",
+                )
+            )
+
+        if (
+            isinstance(node.iter, ast.Call)
+            and isinstance(node.iter.func, ast.Name)
+            and node.iter.func.id == "enumerate"
+        ):
+            errors.append(
+                PIE805(
+                    lineno=node.lineno,
+                    col_offset=node.col_offset,
+                    suggestion=f"use `for {res.var} in enumerate(...)`",
+                )
+            )
+
+
+def pie805_prefer_simple_iterator_generator(
+    node: ast.GeneratorExp | ast.ListComp, errors: list[Error]
+) -> None:
+    for comprehension in node.generators:
+        if (
+            isinstance(comprehension.iter, ast.Call)
+            and isinstance(comprehension.iter.func, ast.Attribute)
+            and comprehension.iter.func.attr == "items"
+            and isinstance(comprehension.target, ast.Tuple)
+            and len(comprehension.target.elts) == 2
+            and is_name_list(comprehension.target.elts)
+            and has_unused_var(comprehension.target.elts)
+        ):
+            res = get_used_var(comprehension.target.elts)
+            if res is None:
+                continue
+            errors.append(
+                PIE805(
+                    lineno=node.lineno,
+                    col_offset=node.col_offset,
+                    suggestion=f"use `for {res.var} in foo.{res.method}()`",
+                )
+            )
+
+
+def PIE805(lineno: int, col_offset: int, suggestion: str) -> Error:
+    return Error(
+        lineno=lineno,
+        col_offset=col_offset,
+        message=f"PIE805: prefer-simple-iterator: {suggestion}",
+    )

flake8_pie/tests/test_pie805_prefer_simple_iterator.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import ast
+
+import pytest
+
+from flake8_pie import Flake8PieCheck
+from flake8_pie.base import Error
+from flake8_pie.pie805_prefer_simple_iterator import PIE805
+from flake8_pie.tests.utils import ex, to_errors
+
+EXAMPLES = [
+    ex(
+        code="""
+for _idx, foo in enumerate(bar):
+    ...
+""",
+        errors=[
+            PIE805(lineno=2, col_offset=0, suggestion="use `for foo in enumerate(...)`")
+        ],
+    ),
+    ex(
+        code="""
+for _key, value in foo.items():
+    ...
+""",
+        errors=[
+            PIE805(lineno=2, col_offset=0, suggestion="use `for value in foo.values()`")
+        ],
+    ),
+    ex(
+        code="""
+for key, _value in foo.items():
+    ...
+""",
+        errors=[
+            PIE805(lineno=2, col_offset=0, suggestion="use `for key in foo.keys()`")
+        ],
+    ),
+    ex(
+        code="""
+fields = [
+    k for k, _v in serialize(user).fields.items() if k != "internal"
+]
+""",
+        errors=[PIE805(lineno=3, col_offset=4, suggestion="use `for k in foo.keys()`")],
+    ),
+    ex(
+        code="""
+users = (
+    User(data)
+    for _id, data in user_map.items()
+    for f, _y in blah.items()
+)
+""",
+        errors=[
+            PIE805(lineno=3, col_offset=4, suggestion="use `for data in foo.values()`"),
+            PIE805(lineno=3, col_offset=4, suggestion="use `for f in foo.keys()`"),
+        ],
+    ),
+    # we need to do usage analysis on the variables to determine if they are actually unused instead of looking at `_`.
+    ex(
+        code="""
+[dict(_name=_name, data=data) for _name, data in users.items()]
+""",
+        errors=[
+            PIE805(lineno=2, col_offset=1, suggestion="use `for data in foo.values()`")
+        ],
+    ),
+    ex(
+        code="""
+for key, value in foo.items():
+    ...
+for idx, foo in enumerate(bar):
+    ...
+[k for k, v in users.items() if v == "guest"]
+""",
+        errors=[],
+    ),
+]
+
+
+@pytest.mark.parametrize("code,errors", EXAMPLES)
+def test_examples(code: str, errors: list[Error]) -> None:
+    expr = ast.parse(code)
+    assert to_errors(Flake8PieCheck(expr, filename="foo.py").run()) == errors

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "flake8-pie"
-version = "0.13.0"
+version = "0.14.0"
 description = "A flake8 extension that implements misc. lints"
 repository = "https://github.com/sbdchd/flake8-pie"
 authors = ["Steve Dignam <[email protected]>"]
@@ -29,9 +29,10 @@ twine = "^1.12"
 pytest-watch = "^4.2"
 pytest = "^4.0"
 ipython = "^7.2"
-mypy = "^0.812.0"
+mypy = "^0.900"
 astpretty = "^2.1"
 isort = "^4.3"
+typing_extensions = "^3.10.0.0"
 
 [tool.poetry.plugins]
 [tool.poetry.plugins."flake8.extension"]