Repository: adamchainz/flake8-comprehensions (python)

Issue:
Sum list vs generator comprehesion detection with startpoint arg
So right now the linter highlights this:
```sum([x + 1 for x in range(30)])``` and for all the right reasons.
However, this isn't detected right now:
```sum([x + 1 for x in range(30)], 0)```

Pull request description:
Fix C407 to catch cases where enumerate and sum have >1 argument
Fixes #178

Patch:
diff --git a/HISTORY.rst b/HISTORY.rst
index 3352127..6df318a 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -16,6 +16,9 @@ Pending Release
 
 * Add dependencies on ``cached-property`` and ``setuptools``.
 
+* Fix false negatives in ``C407`` for cases when ``enumerate`` and ``sum()``
+  are passed more than one argument.
+
 2.2.0 (2019-08-12)
 ------------------
 
diff --git a/src/flake8_comprehensions.py b/src/flake8_comprehensions.py
index 39ae2af..0d67233 100644
--- a/src/flake8_comprehensions.py
+++ b/src/flake8_comprehensions.py
@@ -39,10 +39,10 @@ class ComprehensionChecker:
     def run(self):
         for node in ast.walk(self.tree):
             if isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
-                n_args = len(node.args)
+                num_positional_args = len(node.args)
 
                 if (
-                    n_args == 1
+                    num_positional_args == 1
                     and isinstance(node.args[0], ast.GeneratorExp)
                     and node.func.id in ("list", "set")
                 ):
@@ -55,7 +55,7 @@ class ComprehensionChecker:
                     )
 
                 elif (
-                    n_args == 1
+                    num_positional_args == 1
                     and isinstance(node.args[0], ast.GeneratorExp)
                     and isinstance(node.args[0].elt, ast.Tuple)
                     and len(node.args[0].elt.elts) == 2
@@ -69,7 +69,7 @@ class ComprehensionChecker:
                     )
 
                 elif (
-                    n_args == 1
+                    num_positional_args == 1
                     and isinstance(node.args[0], ast.ListComp)
                     and node.func.id in ("list", "set", "dict")
                 ):
@@ -83,7 +83,7 @@ class ComprehensionChecker:
                         type(self),
                     )
 
-                elif n_args == 1 and (
+                elif num_positional_args == 1 and (
                     isinstance(node.args[0], ast.Tuple)
                     and node.func.id == "tuple"
                     or isinstance(node.args[0], ast.List)
@@ -102,7 +102,7 @@ class ComprehensionChecker:
                     )
 
                 elif (
-                    n_args == 1
+                    num_positional_args == 1
                     and isinstance(node.args[0], (ast.Tuple, ast.List))
                     and node.func.id in ("tuple", "list", "set", "dict")
                 ):
@@ -124,19 +124,36 @@ class ComprehensionChecker:
                     )
 
                 elif (
-                    n_args == 1
+                    num_positional_args == 1
                     and isinstance(node.args[0], ast.ListComp)
                     and node.func.id
                     in (
                         "all",
                         "any",
-                        "enumerate",
                         "frozenset",
+                        "tuple",
+                        # These take 1 positional argument + some keyword arguments
                         "max",
                         "min",
                         "sorted",
+                    )
+                ):
+
+                    yield (
+                        node.lineno,
+                        node.col_offset,
+                        self.messages["C407"].format(func=node.func.id),
+                        type(self),
+                    )
+
+                elif (
+                    num_positional_args in (1, 2)
+                    and isinstance(node.args[0], ast.ListComp)
+                    and node.func.id
+                    in (
+                        # These can take a second positional argument
+                        "enumerate",
                         "sum",
-                        "tuple",
                     )
                 ):
 
@@ -148,7 +165,7 @@ class ComprehensionChecker:
                     )
 
                 elif (
-                    n_args == 0
+                    num_positional_args == 0
                     and not has_star_args(node)
                     and not has_keyword_args(node)
                     and node.func.id in ("tuple", "list", "dict")


Test patch:
diff --git a/tests/test_flake8_comprehensions.py b/tests/test_flake8_comprehensions.py
index 8293551..a86a977 100644
--- a/tests/test_flake8_comprehensions.py
+++ b/tests/test_flake8_comprehensions.py
@@ -332,28 +332,76 @@ def test_C407_sum_fail_1(flake8dir):
     ]
 
 
-def test_C407_max_fail_1(flake8dir):
+def test_C407_max_pass_1(flake8dir):
+    flake8dir.make_example_py("max(x for x in range(10))")
+    result = flake8dir.run_flake8()
+    assert result.out_lines == []
+
+
+def test_C407_max_pass_2(flake8dir):
+    flake8dir.make_example_py("max((x for x in range(10)), key=lambda x: x * 2)")
+    result = flake8dir.run_flake8()
+    assert result.out_lines == []
+
+
+def test_C407_max_pass_3(flake8dir):
     flake8dir.make_example_py(
-        """
-        foo = max([x for x in range(10)])
-    """
+        "max((x for x in range(10)), default=1, key=lambda x: x * 2)"
     )
     result = flake8dir.run_flake8()
+    assert result.out_lines == []
+
+
+def test_C407_max_fail_1(flake8dir):
+    flake8dir.make_example_py("max([x for x in range(10)])")
+    result = flake8dir.run_flake8()
+    assert result.out_lines == [
+        "./example.py:1:1: C407 Unnecessary list comprehension - 'max' can take "
+        + "a generator."
+    ]
+
+
+def test_C407_max_fail_2(flake8dir):
+    flake8dir.make_example_py("max([x for x in range(10)], default=1)")
+    result = flake8dir.run_flake8()
     assert result.out_lines == [
-        "./example.py:1:7: C407 Unnecessary list comprehension - 'max' can take "
+        "./example.py:1:1: C407 Unnecessary list 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()
+    assert result.out_lines == []
+
+
+def test_C407_enumerate_pass_2(flake8dir):
+    flake8dir.make_example_py("enumerate((1 for i in range(10)), 1)")
+    result = flake8dir.run_flake8()
+    assert result.out_lines == []
+
+
+def test_C407_enumerate_pass_3(flake8dir):
+    flake8dir.make_example_py("enumerate((1 for i in range(10)), start=1)")
+    result = flake8dir.run_flake8()
+    assert result.out_lines == []
+
+
 def test_C407_enumerate_fail_1(flake8dir):
-    flake8dir.make_example_py(
-        """
-        foo = enumerate([x for x in range(10)])
-    """
-    )
+    flake8dir.make_example_py("enumerate([x for x in range(10)])")
+    result = flake8dir.run_flake8()
+    assert result.out_lines == [
+        "./example.py:1:1: C407 Unnecessary list comprehension - 'enumerate' "
+        + "can take a generator."
+    ]
+
+
+def test_C407_enumerate_fail_2(flake8dir):
+    flake8dir.make_example_py("enumerate([x for x in range(10)], 1)")
     result = flake8dir.run_flake8()
     assert result.out_lines == [
-        "./example.py:1:7: C407 Unnecessary list comprehension - 'enumerate' "
+        "./example.py:1:1: C407 Unnecessary list comprehension - 'enumerate' "
         + "can take a generator."
     ]