Skip to content

Commit

Permalink
Add C420 rule
Browse files Browse the repository at this point in the history
  • Loading branch information
oliver-sanders committed Aug 4, 2023
1 parent 594fbf8 commit b703e03
Show file tree
Hide file tree
Showing 3 changed files with 86 additions and 0 deletions.
12 changes: 12 additions & 0 deletions README.rst
Original file line number Diff line number Diff line change
Expand Up @@ -227,3 +227,15 @@ For example:

* Rewrite ``all([condition(x) for x in iterable])`` as ``all(condition(x) for x in iterable)``
* Rewrite ``any([condition(x) for x in iterable])`` as ``any(condition(x) for x in iterable)``

C420 Unnecessary list comprehension in for loop rewrite as a generator.
-----------------------------------------------------------------------

It's unnecessary to iterate over a list comprehension in a for loop.

This prevents short circuiting (e.g. the for loop may contain a ``break`` statement) and creates an unnecessary intermediate list.

Use a generator expression instead:

* Rewrite ``for x in [func(y) for y in iterable]:`` as ``for x in (func(y) for y in iterable):``
* Rewrite ``async for x in [func(y) for y in async_gen]:`` as ``async for x in (func(y) for y in async_gen):``
21 changes: 21 additions & 0 deletions src/flake8_comprehensions/__init__.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

import ast
from copy import deepcopy
from importlib.metadata import version
from typing import Any
from typing import Generator
Expand Down Expand Up @@ -46,6 +47,10 @@ def __init__(self, tree: ast.AST) -> None:
"C419 Unnecessary list comprehension passed to {func}() prevents "
+ "short-circuiting - rewrite as a generator."
),
"C420": (
"C420 Unnecessary list comprehension in for loop "
+ "rewrite as a generator: {fix}"
),
}

def run(self) -> Generator[tuple[int, int, str, type[Any]], None, None]:
Expand Down Expand Up @@ -361,6 +366,22 @@ def run(self) -> Generator[tuple[int, int, str, type[Any]], None, None]:
self.messages["C416"].format(type=comp_type[node.__class__]),
type(self),
)
elif isinstance(node, (ast.AsyncFor, ast.For)) and isinstance(
node.iter, ast.ListComp
):
# create a new node to fix the code:
_node = deepcopy(node)
# clear the body of the for loop
_node.body = []
# replace the list comprehension with a generator expression
_node.iter = ast.GeneratorExp(node.iter.elt, node.iter.generators)

yield (
node.lineno,
node.col_offset,
self.messages["C420"].format(fix=ast.unparse(_node)),
type(self),
)


def has_star_args(call_node: ast.Call) -> bool:
Expand Down
53 changes: 53 additions & 0 deletions tests/test_flake8_comprehensions.py
Original file line number Diff line number Diff line change
Expand Up @@ -963,3 +963,56 @@ def test_C419_fail(code, failures, flake8_path):
(flake8_path / "example.py").write_text(dedent(code))
result = flake8_path.run_flake8()
assert result.out_lines == failures


@pytest.mark.parametrize(
"code",
[
"for x in (y ** 2 for y in range(10)): pass",
"async for x in (fcn(y) for y in async_gen): pass",
"for x in y: pass",
"for x in y: {fcn(y) for y in range(10)}",
"for x, y in y: {str(y): y for y in range(10)}.items()",
],
)
def test_C420_pass(code, flake8_path):
(flake8_path / "example.py").write_text(dedent(code))
result = flake8_path.run_flake8()
assert result.out_lines == []


@pytest.mark.parametrize(
"code,failures",
[
(
# it should detect unnecessary list comprehensions in for loops
"for x in [y ** 2 for y in range(10)]: pass",
[
"./example.py:1:1: C420 Unnecessary list comprehension in for loop "
+ "rewrite as a generator: for x in (y ** 2 for y in range(10)):"
],
),
(
# it should also run other rules over the comprenension
"for x in [y for y in range(10)]: pass",
[
"./example.py:1:1: C420 Unnecessary list comprehension in for loop "
+ "rewrite as a generator: for x in (y for y in range(10)):",
"./example.py:1:10: C416 Unnecessary list comprehension - rewrite using "
"list().",
],
),
(
# it should work for async for loops too
"async for x in [fcn(y) for y in async_gen]: pass",
[
"./example.py:1:1: C420 Unnecessary list comprehension in for loop "
+ "rewrite as a generator: async for x in (fcn(y) for y in async_gen):",
],
),
],
)
def test_C420_fail(code, failures, flake8_path):
(flake8_path / "example.py").write_text(dedent(code))
result = flake8_path.run_flake8()
assert result.out_lines == failures

0 comments on commit b703e03

Please sign in to comment.