From 5c3998187e39730fd84088d71e555c9cddbc9461 Mon Sep 17 00:00:00 2001
From: Martin Thoma <info@martin-thoma.de>
Date: Tue, 17 Nov 2020 23:48:27 +0100
Subject: [PATCH] ENH: Add SIM113 (enumerate in for-loops)

Closes issue #18
---
 flake8_simplify.py     | 48 ++++++++++++++++++++++++++++++++++++++++++
 tests/test_simplify.py |  8 +++++++
 2 files changed, 56 insertions(+)

diff --git a/flake8_simplify.py b/flake8_simplify.py
index 14b8a41..e116a08 100644
--- a/flake8_simplify.py
+++ b/flake8_simplify.py
@@ -33,6 +33,7 @@
 SIM110 = "SIM110 Use 'return any({check} for {target} in {iterable})'"
 SIM111 = "SIM111 Use 'return all({check} for {target} in {iterable})'"
 SIM112 = "SIM112 Use '{expected}' instead of '{original}'"
+SIM113 = "SIM113 Use enumerate instead of '{variable}'"
 SIM201 = "SIM201 Use '{left} != {right}' instead of 'not {left} == {right}'"
 SIM202 = "SIM202 Use '{left} == {right}' instead of 'not {left} != {right}'"
 SIM203 = "SIM203 Use '{a} not in {b}' instead of 'not {a} in {b}'"
@@ -650,6 +651,52 @@ def _get_sim112(node: ast.Expr) -> List[Tuple[int, int, str]]:
     return errors
 
 
+def _get_sim113(node: ast.For) -> List[Tuple[int, int, str]]:
+    """
+    Find loops in which "enumerate" should be used.
+
+        For(
+            target=Name(id='el', ctx=Store()),
+            iter=Name(id='iterable', ctx=Load()),
+            body=[
+                Expr(
+                    value=Constant(value=Ellipsis, kind=None),
+                ),
+                AugAssign(
+                    target=Name(id='idx', ctx=Store()),
+                    op=Add(),
+                    value=Constant(value=1, kind=None),
+                ),
+            ],
+            orelse=[],
+            type_comment=None,
+        ),
+    """
+    errors: List[Tuple[int, int, str]] = []
+    variable_candidates = []
+    for expression in node.body:
+        if isinstance(expression, ast.AugAssign) and is_constant_increase(
+            expression
+        ):
+            variable_candidates.append(expression.target)
+
+    for candidate in variable_candidates:
+        errors.append(
+            (
+                candidate.lineno,
+                candidate.col_offset,
+                SIM113.format(variable=to_source(candidate)),
+            )
+        )
+    return errors
+
+
+def is_constant_increase(expr: ast.AugAssign) -> bool:
+    return isinstance(expr.op, ast.Add) and isinstance(
+        expr.value, ast.Constant
+    )
+
+
 def _get_sim201(node: ast.UnaryOp) -> List[Tuple[int, int, str]]:
     """
     Get a list of all calls where an unary 'not' is used for an equality.
@@ -1063,6 +1110,7 @@ def visit_If(self, node: ast.If) -> None:
     def visit_For(self, node: ast.For) -> None:
         self.errors += _get_sim104(node)
         self.errors += _get_sim110_sim111(node)
+        self.errors += _get_sim113(node)
         self.generic_visit(node)
 
     def visit_Try(self, node: ast.Try) -> None:
diff --git a/tests/test_simplify.py b/tests/test_simplify.py
index 60d1979..5ca66e2 100644
--- a/tests/test_simplify.py
+++ b/tests/test_simplify.py
@@ -280,6 +280,14 @@ def test_sim_112_get_with_default():
     }
 
 
+def test_sim_113():
+    ret = _results(
+        """for el in iterable:
+    idx += 1"""
+    )
+    assert ret == {"2:4 SIM113 Use enumerate instead of 'idx'"}
+
+
 def test_sim_201():
     ret = _results("not a == b")
     assert ret == {"1:0 SIM201 Use 'a != b' instead of 'not a == b'"}