Skip to content

Commit

Permalink
Refactor implementation slightly
Browse files Browse the repository at this point in the history
  • Loading branch information
adamchainz committed May 19, 2022
1 parent 293b116 commit acfba82
Showing 1 changed file with 16 additions and 14 deletions.
30 changes: 16 additions & 14 deletions src/flake8_comprehensions/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -273,23 +273,25 @@ def run(self) -> Generator[tuple[int, int, str, type[Any]], None, None]:
map_call = node.args[0]
visited_map_calls.add(map_call)

rewriteable = True
if node.func.id == "dict":
# For a dict comprehension to be able to be rewritten as
# a comprehension, the lambda body should be a 2-tuple.
# For the generator expression to be rewriteable as a
# dict comprehension, its lambda must return a 2-tuple.
lambda_node = node.args[0].args[0]
if not (
isinstance(lambda_node.body, (ast.List, ast.Tuple))
and len(lambda_node.body.elts) == 2
if (
not isinstance(lambda_node.body, (ast.List, ast.Tuple))
or len(lambda_node.body.elts) != 2
):
continue

comprehension_type = f"{node.func.id} comprehension"
yield (
node.lineno,
node.col_offset,
self.messages["C417"].format(comp=comprehension_type),
type(self),
)
rewriteable = False

if rewriteable:
comprehension_type = f"{node.func.id} comprehension"
yield (
node.lineno,
node.col_offset,
self.messages["C417"].format(comp=comprehension_type),
type(self),
)

elif isinstance(node, (ast.ListComp, ast.SetComp)):
if (
Expand Down

0 comments on commit acfba82

Please sign in to comment.