Add C417 - Unnecessary use of map (#409)

Co-authored-by: Adam Johnson <[email protected]>

HISTORY.rst

@@ -2,6 +2,10 @@
 History
 =======
 
+* Add rule C417 which recommends rewriting use of ``map()`` with ``lambda`` to an equivalent generator expression or comprehension.
+
+ Thanks to Tushar Sadhwani in `PR #409 <https:/adamchainz/flake8-comprehensions/pull/409>`__.
+
 3.9.0 (2022-05-11)
 ------------------
 

README.rst

@@ -159,3 +159,16 @@ For example:
 
 * Rewrite ``[x for x in iterable]`` as ``list(iterable)``
 * Rewrite ``{x for x in iterable}`` as ``set(iterable)``
+
+C417: Unnecessary ``map`` usage - rewrite using a generator expression/``<list/set/dict>`` comprehension.
+---------------------------------------------------------------------------------------------------------
+
+``map(func, iterable)`` has great performance when ``func`` is a built-in function, and it makes sense if your function already has a name.
+But if your func is a ``lambda``, it’s faster to use a generator expression or a comprehension, as it avoids the function call overhead.
+For example:
+
+* Rewrite ``map(lambda x: x + 1, iterable)`` to ``(x + 1 for x in iterable)``
+* Rewrite ``map(lambda item: get_id(item), items)`` to ``(get_id(item) for item in items)``
+* Rewrite ``list(map(lambda num: num * 2, nums))`` to ``[num * 2 for num in nums]``
+* Rewrite ``set(map(lambda num: num % 2 == 0, nums))`` to ``{num % 2 == 0 for num in nums}``
+* Rewrite ``dict(map(lambda v: (v, v ** 2), values))`` to ``{v : v ** 2 for v in values}``

src/flake8_comprehensions/__init__.py

@@ -41,9 +41,13 @@ def __init__(self, tree: ast.AST) -> None:
  "C414": "C414 Unnecessary {inner} call within {outer}().",
  "C415": "C415 Unnecessary subscript reversal of iterable within {func}().",
  "C416": "C416 Unnecessary {type} comprehension - rewrite using {type}().",
+ "C417": "C417 Unnecessary use of map - use a {comp} instead.",
  }
 
  def run(self) -> Generator[tuple[int, int, str, type[Any]], None, None]:
+ # Stores previously seen map() nodes, to avoid raising C417 on it twice.
+ visited_map_calls: set[ast.Call] = set()
+
  for node in ast.walk(self.tree):
  if isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
  num_positional_args = len(node.args)
@@ -243,6 +247,52 @@ def run(self) -> Generator[tuple[int, int, str, type[Any]], None, None]:
  type(self),
  )
 
+ elif (
+ node.func.id == "map"
+ and node not in visited_map_calls
+ and len(node.args) == 2
+ and isinstance(node.args[0], ast.Lambda)
+ ):
+ yield (
+ node.lineno,
+ node.col_offset,
+ self.messages["C417"].format(comp="generator expression"),
+ type(self),
+ )
+
+ elif (
+ node.func.id in ("list", "set", "dict")
+ and len(node.args) == 1
+ and isinstance(node.args[0], ast.Call)
+ and isinstance(node.args[0].func, ast.Name)
+ and node.args[0].func.id == "map"
+ and len(node.args[0].args) == 2
+ and isinstance(node.args[0].args[0], ast.Lambda)
+ ):
+ # To avoid raising C417 on the map() call inside the list/set/dict.
+ map_call = node.args[0]
+ visited_map_calls.add(map_call)
+
+ rewriteable = True
+ if node.func.id == "dict":
+ # 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))
+ or len(lambda_node.body.elts) != 2
+ ):
+ 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 (
  len(node.generators) == 1

tests/test_flake8_comprehensions.py

@@ -773,3 +773,60 @@ def test_C416_fail_2_set(flake8_path):
  assert result.out_lines == [
  "./example.py:1:1: C416 Unnecessary set comprehension - rewrite using set().",
  ]
+
+
+# C417
+
+
+@pytest.mark.parametrize(
+ "code",
+ [
+ "map()",
+ "map(str, numbers)",
+ "list(map())",
+ "list(map(str, numbers))",
+ "set(map(f, items))",
+ "dict(map(enumerate, values))",
+ "dict(map(lambda v: data[v], values))",
+ ],
+)
+def test_C417_pass(code, flake8_path):
+ (flake8_path / "example.py").write_text(code)
+ result = flake8_path.run_flake8()
+ assert result.out_lines == []
+
+
+def test_C417_fail_1_generator_expression(flake8_path):
+ (flake8_path / "example.py").write_text("map(lambda x: x * 2, iterable)")
+ result = flake8_path.run_flake8()
+ assert result.out_lines == [
+ "./example.py:1:1: C417 Unnecessary use of map - "
+ + "use a generator expression instead.",
+ ]
+
+
+def test_C417_fail_2_list_comprehension(flake8_path):
+ (flake8_path / "example.py").write_text("list(map(lambda x: x * 2, iterable))")
+ result = flake8_path.run_flake8()
+ assert result.out_lines == [
+ "./example.py:1:1: C417 Unnecessary use of map - "
+ + "use a list comprehension instead.",
+ ]
+
+
+def test_C417_fail_3_set_comprehension(flake8_path):
+ (flake8_path / "example.py").write_text("set(map(lambda num: num % 2 == 0, nums))")
+ result = flake8_path.run_flake8()
+ assert result.out_lines == [
+ "./example.py:1:1: C417 Unnecessary use of map - "
+ + "use a set comprehension instead.",
+ ]
+
+
+def test_C417_fail_4_dict_comprehension(flake8_path):
+ (flake8_path / "example.py").write_text("dict(map(lambda v: (v, v ** 2), values))")
+ result = flake8_path.run_flake8()
+ assert result.out_lines == [
+ "./example.py:1:1: C417 Unnecessary use of map - "
+ "use a dict comprehension instead.",
+ ]