Skip to content

Commit

Permalink
ENH: Add SIM113 (enumerate in for-loops)
Browse files Browse the repository at this point in the history
Closes issue #18
  • Loading branch information
MartinThoma committed Nov 17, 2020
1 parent b63f127 commit 5c39981
Show file tree
Hide file tree
Showing 2 changed files with 56 additions and 0 deletions.
48 changes: 48 additions & 0 deletions flake8_simplify.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}'"
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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:
Expand Down
8 changes: 8 additions & 0 deletions tests/test_simplify.py
Original file line number Diff line number Diff line change
Expand Up @@ -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'"}
Expand Down

0 comments on commit 5c39981

Please sign in to comment.