Allow additional arguments for sum and max comprehensions (#12364)

## Summary These can have other arguments, so it seems wrong to gate on single argument here. Closes #12358.
astral-sh · Jul 18, 2024 · 9b9d701 · 9b9d701
1 parent 648cca1
commit 9b9d701
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 17 deletions.
diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_comprehensions/C419_1.py b/crates/ruff_linter/resources/test/fixtures/flake8_comprehensions/C419_1.py
@@ -1,8 +1,10 @@
 sum([x.val for x in bar])
 min([x.val for x in bar])
 max([x.val for x in bar])
+sum([x.val for x in bar], 0)
 
 # Ok
 sum(x.val for x in bar)
 min(x.val for x in bar)
 max(x.val for x in bar)
+sum(x.val for x in bar, 0)
diff --git a/crates/ruff_linter/src/rules/flake8_comprehensions/fixes.rs b/crates/ruff_linter/src/rules/flake8_comprehensions/fixes.rs
@@ -3,11 +3,11 @@ use std::iter;
 use anyhow::{bail, Result};
 use itertools::Itertools;
 use libcst_native::{
-    Arg, AssignEqual, AssignTargetExpression, Call, Comment, CompFor, Dict, DictComp, DictElement,
-    Element, EmptyLine, Expression, GeneratorExp, LeftCurlyBrace, LeftParen, LeftSquareBracket,
-    ListComp, Name, ParenthesizableWhitespace, ParenthesizedNode, ParenthesizedWhitespace,
-    RightCurlyBrace, RightParen, RightSquareBracket, SetComp, SimpleString, SimpleWhitespace,
-    TrailingWhitespace, Tuple,
+    Arg, AssignEqual, AssignTargetExpression, Call, Comma, Comment, CompFor, Dict, DictComp,
+    DictElement, Element, EmptyLine, Expression, GeneratorExp, LeftCurlyBrace, LeftParen,
+    LeftSquareBracket, ListComp, Name, ParenthesizableWhitespace, ParenthesizedNode,
+    ParenthesizedWhitespace, RightCurlyBrace, RightParen, RightSquareBracket, SetComp,
+    SimpleString, SimpleWhitespace, TrailingWhitespace, Tuple,
 };
 
 use ruff_diagnostics::{Edit, Fix};
@@ -937,7 +937,10 @@ pub(crate) fn fix_unnecessary_comprehension_in_call(
     let whitespace_after_arg = match &call.args[0].comma {
         Some(comma) => {
             let whitespace_after_comma = comma.whitespace_after.clone();
-            call.args[0].comma = None;
+            call.args[0].comma = Some(Comma {
+                whitespace_after: ParenthesizableWhitespace::default(),
+                ..comma.clone()
+            });
             whitespace_after_comma
         }
         _ => call.args[0].whitespace_after_arg.clone(),

diff --git a/...es/ruff_linter/src/rules/flake8_comprehensions/rules/unnecessary_comprehension_in_call.rs b/...es/ruff_linter/src/rules/flake8_comprehensions/rules/unnecessary_comprehension_in_call.rs
@@ -90,7 +90,7 @@ pub(crate) fn unnecessary_comprehension_in_call(
     if !keywords.is_empty() {
         return;
     }
-    let [arg] = args else {
+    let Some(arg) = args.first() else {
         return;
     };
     let (Expr::ListComp(ast::ExprListComp { elt, .. })

diff --git a/...prehensions/snapshots/ruff_linter__rules__flake8_comprehensions__tests__C419_C419.py.snap b/...prehensions/snapshots/ruff_linter__rules__flake8_comprehensions__tests__C419_C419.py.snap
@@ -51,7 +51,7 @@ C419.py:4:5: C419 [*] Unnecessary list comprehension
 2 2 | all([x.id for x in bar])
 3 3 | any(  # first comment
 4   |-    [x.id for x in bar],  # second comment
-  4 |+    x.id for x in bar  # second comment
+  4 |+    x.id for x in bar,  # second comment
 5 5 | )  # third comment
 6 6 | all(  # first comment
 7 7 |     [x.id for x in bar],  # second comment
@@ -72,7 +72,7 @@ C419.py:7:5: C419 [*] Unnecessary list comprehension
 5 5 | )  # third comment
 6 6 | all(  # first comment
 7   |-    [x.id for x in bar],  # second comment
-  7 |+    x.id for x in bar  # second comment
+  7 |+    x.id for x in bar,  # second comment
 8 8 | )  # third comment
 9 9 | any({x.id for x in bar})
 10 10 | 

diff --git a/.../snapshots/ruff_linter__rules__flake8_comprehensions__tests__preview__C419_C419_1.py.snap b/.../snapshots/ruff_linter__rules__flake8_comprehensions__tests__preview__C419_C419_1.py.snap
@@ -15,14 +15,15 @@ C419_1.py:1:5: C419 [*] Unnecessary list comprehension
   1 |+sum(x.val for x in bar)
 2 2 | min([x.val for x in bar])
 3 3 | max([x.val for x in bar])
-4 4 | 
+4 4 | sum([x.val for x in bar], 0)
 
 C419_1.py:2:5: C419 [*] Unnecessary list comprehension
   |
 1 | sum([x.val for x in bar])
 2 | min([x.val for x in bar])
   |     ^^^^^^^^^^^^^^^^^^^^ C419
 3 | max([x.val for x in bar])
+4 | sum([x.val for x in bar], 0)
   |
   = help: Remove unnecessary list comprehension
 
@@ -31,17 +32,16 @@ C419_1.py:2:5: C419 [*] Unnecessary list comprehension
 2   |-min([x.val for x in bar])
   2 |+min(x.val for x in bar)
 3 3 | max([x.val for x in bar])
-4 4 | 
-5 5 | # Ok
+4 4 | sum([x.val for x in bar], 0)
+5 5 | 
 
 C419_1.py:3:5: C419 [*] Unnecessary list comprehension
   |
 1 | sum([x.val for x in bar])
 2 | min([x.val for x in bar])
 3 | max([x.val for x in bar])
   |     ^^^^^^^^^^^^^^^^^^^^ C419
-4 | 
-5 | # Ok
+4 | sum([x.val for x in bar], 0)
   |
   = help: Remove unnecessary list comprehension
 
@@ -50,6 +50,27 @@ C419_1.py:3:5: C419 [*] Unnecessary list comprehension
 2 2 | min([x.val for x in bar])
 3   |-max([x.val for x in bar])
   3 |+max(x.val for x in bar)
-4 4 | 
-5 5 | # Ok
-6 6 | sum(x.val for x in bar)
+4 4 | sum([x.val for x in bar], 0)
+5 5 | 
+6 6 | # Ok
+
+C419_1.py:4:5: C419 [*] Unnecessary list comprehension
+  |
+2 | min([x.val for x in bar])
+3 | max([x.val for x in bar])
+4 | sum([x.val for x in bar], 0)
+  |     ^^^^^^^^^^^^^^^^^^^^ C419
+5 | 
+6 | # Ok
+  |
+  = help: Remove unnecessary list comprehension
+
+ℹ Unsafe fix
+1 1 | sum([x.val for x in bar])
+2 2 | min([x.val for x in bar])
+3 3 | max([x.val for x in bar])
+4   |-sum([x.val for x in bar], 0)
+  4 |+sum(x.val for x in bar, 0)
+5 5 | 
+6 6 | # Ok
+7 7 | sum(x.val for x in bar)