Repository: adamchainz/flake8-comprehensions
Pull request #215 (2020-01-20 09:25:28)

Check for dict and set comprehensions in rules C407 and C412
Mentioned in #212.


Patch:
diff --git a/HISTORY.rst b/HISTORY.rst
index f906c04..a9177dc 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -3,6 +3,7 @@ History
 =======
 
 * Add ``filter`` and ``map`` to rule C407.
+* Check for dict and set comprehensions in rules C407 and C412.
 
 3.1.4 (2019-11-20)
 ------------------
diff --git a/README.rst b/README.rst
index 674dcb7..59390fa 100644
--- a/README.rst
+++ b/README.rst
@@ -47,12 +47,12 @@ C403 Unnecessary list comprehension - rewrite as a set comprehension.
 C404 Unnecessary list comprehension - rewrite as a dict comprehension.
 C405 Unnecessary (list/tuple) literal - rewrite as a set literal.
 C406 Unnecessary (list/tuple) literal - rewrite as a dict literal.
-C407 Unnecessary list comprehension - '<builtin>' can take a generator.
+C407 Unnecessary (dict/list/set) comprehension - '<builtin>' can take a generator.
 C408 Unnecessary (dict/list/tuple) call - rewrite as a literal.
 C409 Unnecessary (list/tuple) passed to tuple() - (remove the outer call to tuple()/rewrite as a tuple literal).
 C410 Unnecessary (list/tuple) passed to list() - (remove the outer call to list()/rewrite as a list literal).
 C411 Unnecessary list call - remove the outer call to list().
-C412 Unnecessary list comprehension - 'in' can take a generator.
+C412 Unnecessary (dict/list/set) comprehension - 'in' can take a generator.
 C413 Unnecessary list call around sorted().
 C413 Unnecessary reversed call around sorted() - (use sorted(..., reverse=(True/False))/toggle reverse argument to sorted()).
 C414 Unnecessary (list/reversed/set/sorted/tuple) call within list/set/sorted/tuple().
@@ -100,16 +100,16 @@ For example:
 * Rewrite ``dict([(1, 2)])`` or ``dict(((1, 2),))`` as ``{1: 2}``
 * Rewrite ``dict([])`` as ``{}``
 
-C407: Unnecessary list comprehension - '<builtin>' can take a generator
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+C407: Unnecessary (dict/list/set) comprehension - '<builtin>' can take a generator
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 It's unnecessary to pass a list comprehension to some builtins that can take
 generators instead. For example:
 
-* Rewrite ``sum([x ** 2 for x in range(10)])`` as
+* Rewrite ``sum({x ** 2 for x in range(10)})`` as
   ``sum(x ** 2 for x in range(10))``
 * Rewrite ``all([foo.bar for foo in foos])`` as
-  ``all(foo.bar for foo in foos)``
+  ``all(foo.bare for foo in foos)``
 * Rewrite ``filter(lambda x: x % 2 == 0, [x ** 3 for x in range(10)])`` as
   ``filter(lambda x: x % 2 == 0, (x ** 3 for x in range(10)))``
 
@@ -146,13 +146,14 @@ equivalent without it. For example:
 
 * Rewrite ``list([f(x) for x in foo])`` as ``[f(x) for x in foo]``
 
-C412: Unnecessary list comprehension - 'in' can take a generator.
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+C412: Unnecessary (dict/list/set) comprehension - 'in' can take a generator.
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-It's unnecessary to pass a list comprehension to 'in' that can take a
+It's unnecessary to pass a dict/list/set comprehension to 'in' that can take a
 generator instead. For example:
 
 * Rewrite ``y in [f(x) for x in foo]`` as ``y in (f(x) for x in foo)``
+* Rewrite ``y in {x ** 2 for x in foo}`` as ``y in (f(x) for x in foo)``
 
 C413: Unnecessary list/reversed call around sorted().
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/src/flake8_comprehensions.py b/src/flake8_comprehensions.py
index 4f4a114..dc3522e 100644
--- a/src/flake8_comprehensions.py
+++ b/src/flake8_comprehensions.py
@@ -28,12 +28,14 @@ class ComprehensionChecker:
         ),
         "C405": "C405 Unnecessary {type} literal - ",
         "C406": "C406 Unnecessary {type} literal - ",
-        "C407": "C407 Unnecessary list comprehension - '{func}' can take a generator.",
+        "C407": (
+            "C407 Unnecessary {type} comprehension - '{func}' can take a generator."
+        ),
         "C408": "C408 Unnecessary {type} call - rewrite as a literal.",
         "C409": "C409 Unnecessary {type} passed to tuple() - ",
         "C410": "C410 Unnecessary {type} passed to list() - ",
         "C411": "C411 Unnecessary list call - remove the outer call to list().",
-        "C412": "C412 Unnecessary list comprehension - 'in' can take a generator.",
+        "C412": "C412 Unnecessary {type} comprehension - 'in' can take a generator.",
         "C413": "C413 Unnecessary {outer} call around {inner}(){remediation}.",
         "C414": "C414 Unnecessary {inner} call within {outer}().",
         "C415": "C415 Unnecessary subscript reversal of iterable within {func}().",
@@ -136,13 +138,17 @@ class ComprehensionChecker:
                         node.func.id
                         in ("all", "any", "frozenset", "tuple", "max", "min", "sorted",)
                     )
-                    and isinstance(node.args[0], ast.ListComp)
+                    and isinstance(
+                        node.args[0], (ast.DictComp, ast.ListComp, ast.SetComp)
+                    )
                 ):
 
                     yield (
                         node.lineno,
                         node.col_offset,
-                        self.messages["C407"].format(func=node.func.id),
+                        self.messages["C407"].format(
+                            type=comp_type[node.args[0].__class__], func=node.func.id
+                        ),
                         type(self),
                     )
 
@@ -150,42 +156,49 @@ class ComprehensionChecker:
                     num_positional_args in (1, 2)
                     # These can take a second positional argument
                     and (node.func.id in ("enumerate", "sum",))
-                    and isinstance(node.args[0], ast.ListComp)
+                    and isinstance(
+                        node.args[0], (ast.DictComp, ast.ListComp, ast.SetComp)
+                    )
                 ):
 
                     yield (
                         node.lineno,
                         node.col_offset,
-                        self.messages["C407"].format(func=node.func.id),
+                        self.messages["C407"].format(
+                            type=comp_type[node.args[0].__class__], func=node.func.id
+                        ),
                         type(self),
                     )
 
                 elif (
                     num_positional_args == 2
                     and node.func.id == "filter"
-                    and isinstance(node.args[1], ast.ListComp)
-                ):
-
-                    yield (
-                        node.lineno,
-                        node.col_offset,
-                        self.messages["C407"].format(func=node.func.id),
-                        type(self),
+                    and isinstance(
+                        node.args[1], (ast.DictComp, ast.ListComp, ast.SetComp)
                     )
-
-                elif (
-                    num_positional_args >= 2
-                    and node.func.id == "map"
-                    and any(isinstance(arg, ast.ListComp) for arg in node.args[1:])
                 ):
 
                     yield (
                         node.lineno,
                         node.col_offset,
-                        self.messages["C407"].format(func=node.func.id),
+                        self.messages["C407"].format(
+                            type=comp_type[node.args[1].__class__], func=node.func.id
+                        ),
                         type(self),
                     )
 
+                elif num_positional_args >= 2 and node.func.id == "map":
+                    for arg in node.args[1:]:
+                        if isinstance(arg, (ast.DictComp, ast.ListComp, ast.SetComp)):
+                            yield (
+                                node.lineno,
+                                node.col_offset,
+                                self.messages["C407"].format(
+                                    type=comp_type[arg.__class__], func=node.func.id
+                                ),
+                                type(self),
+                            )
+
                 elif (
                     num_positional_args == 0
                     and not has_star_args(node)
@@ -289,12 +302,16 @@ class ComprehensionChecker:
                     len(node.ops) == 1
                     and isinstance(node.ops[0], ast.In)
                     and len(node.comparators) == 1
-                    and isinstance(node.comparators[0], ast.ListComp)
+                    and isinstance(
+                        node.comparators[0], (ast.DictComp, ast.ListComp, ast.SetComp)
+                    )
                 ):
                     yield (
                         node.lineno,
                         node.col_offset,
-                        self.messages["C412"],
+                        self.messages["C412"].format(
+                            type=comp_type[node.comparators[0].__class__]
+                        ),
                         type(self),
                     )
 
@@ -309,11 +326,11 @@ class ComprehensionChecker:
                         and node.elt.id == node.generators[0].target.id
                     )
                 ):
-                    lookup = {ast.ListComp: "list", ast.SetComp: "set"}
+
                     yield (
                         node.lineno,
                         node.col_offset,
-                        self.messages["C416"].format(type=lookup[node.__class__]),
+                        self.messages["C416"].format(type=comp_type[node.__class__]),
                         type(self),
                     )
 
@@ -326,6 +343,13 @@ def has_keyword_args(call_node):
     return any(k.arg is None for k in call_node.keywords)
 
 
+comp_type = {
+    ast.DictComp: "dict",
+    ast.ListComp: "list",
+    ast.SetComp: "set",
+}
+
+
 if sys.version_info >= (3, 6):
 
     def is_async_generator(node):


Test patch:
diff --git a/tests/test_flake8_comprehensions.py b/tests/test_flake8_comprehensions.py
index 72dcae7..14589e6 100644
--- a/tests/test_flake8_comprehensions.py
+++ b/tests/test_flake8_comprehensions.py
@@ -374,6 +374,19 @@ def test_C407_sum_fail_1(flake8dir):
     ]
 
 
+def test_C407_sum_fail_2(flake8dir):
+    flake8dir.make_example_py(
+        """
+        foo = sum({x + 1 for x in range(10)})
+    """
+    )
+    result = flake8dir.run_flake8()
+    assert result.out_lines == [
+        "./example.py:1:7: C407 Unnecessary set comprehension - 'sum' can take "
+        + "a generator."
+    ]
+
+
 def test_C407_filter_pass_1(flake8dir):
     flake8dir.make_example_py(
         "filter(lambda x: x % 2 == 0, (x + 1 for x in range(10)))"
@@ -393,6 +406,17 @@ def test_C407_filter_fail_1(flake8dir):
     ]
 
 
+def test_C407_filter_fail_2(flake8dir):
+    flake8dir.make_example_py(
+        "filter(lambda x: x % 2 == 0, {x + 1 for x in range(10)})"
+    )
+    result = flake8dir.run_flake8()
+    assert result.out_lines == [
+        "./example.py:1:1: C407 Unnecessary set comprehension - 'filter' can take "
+        + "a generator."
+    ]
+
+
 def test_C407_map_pass_1(flake8dir):
     flake8dir.make_example_py("map(lambda x: x + 1, (x + 1 for x in range(10)))")
     result = flake8dir.run_flake8()
@@ -419,6 +443,15 @@ def test_C407_map_fail_2(flake8dir):
     ]
 
 
+def test_C407_map_fail_3(flake8dir):
+    flake8dir.make_example_py("map(lambda x: x + 1, {x + 1 for x in range(10)})")
+    result = flake8dir.run_flake8()
+    assert result.out_lines == [
+        "./example.py:1:1: C407 Unnecessary set comprehension - 'map' can take "
+        + "a generator."
+    ]
+
+
 def test_C407_max_pass_1(flake8dir):
     flake8dir.make_example_py("max(x + 1 for x in range(10))")
     result = flake8dir.run_flake8()
@@ -457,6 +490,15 @@ def test_C407_max_fail_2(flake8dir):
     ]
 
 
+def test_C407_max_fail_3(flake8dir):
+    flake8dir.make_example_py("max({x + 1 for x in range(10)})")
+    result = flake8dir.run_flake8()
+    assert result.out_lines == [
+        "./example.py:1:1: C407 Unnecessary set comprehension - 'max' can take "
+        + "a generator."
+    ]
+
+
 def test_C407_enumerate_pass_1(flake8dir):
     flake8dir.make_example_py("enumerate(1 for i in range(10))")
     result = flake8dir.run_flake8()
@@ -810,6 +852,32 @@ def test_C412_fail_1(flake8dir):
     ]
 
 
+def test_C412_fail_2(flake8dir):
+    flake8dir.make_example_py(
+        """
+        10 in {x + 1 for x in range(10)}
+    """
+    )
+    result = flake8dir.run_flake8()
+    assert result.out_lines == [
+        "./example.py:1:1: C412 Unnecessary set comprehension - 'in' can "
+        + "take a generator."
+    ]
+
+
+def test_C412_fail_3(flake8dir):
+    flake8dir.make_example_py(
+        """
+        10 in {x: x + 1 for x in range(10)}
+    """
+    )
+    result = flake8dir.run_flake8()
+    assert result.out_lines == [
+        "./example.py:1:1: C412 Unnecessary dict comprehension - 'in' can "
+        + "take a generator."
+    ]
+
+
 # C413