Respect trailing whitespace in comprehension fixes (#554)

resources/test/fixtures/C400.py

@@ -1 +1,4 @@
 x = list(x for x in range(3))
+x = list(
+    x for x in range(3)
+)

resources/test/fixtures/C401.py

@@ -1 +1,4 @@
 x = set(x for x in range(3))
+x = set(
+    x for x in range(3)
+)

resources/test/fixtures/C403.py

@@ -1 +1,4 @@
 s = set([x for x in range(3)])
+s = set(
+    [x for x in range(3)]
+)

src/flake8_comprehensions/fixes.rs

@@ -37,8 +37,13 @@ pub fn fix_unnecessary_generator_list(locator: &SourceCodeLocator, expr: &Expr)
     } else {
         return Err(anyhow::anyhow!("Expected node to be: Expression::Call."));
     };
-    let arg = if let Some(Arg { value, .. }) = call.args.first_mut() {
+    let (arg, whitespace_after_arg) = if let Some(Arg {
-        value
+        value,
+        whitespace_after_arg,
+        ..
+    }) = call.args.first_mut()
+    {
+        (value, whitespace_after_arg)
     } else {
         return Err(anyhow::anyhow!("Expected node to be: Arg."));
     };
@@ -57,7 +62,7 @@ pub fn fix_unnecessary_generator_list(locator: &SourceCodeLocator, expr: &Expr)
             whitespace_after: call.whitespace_before_args.clone(),
         },
         rbracket: RightSquareBracket {
-            whitespace_before: call.whitespace_after_func.clone(),
+            whitespace_before: whitespace_after_arg.clone(),
         },
         lpar: generator_exp.lpar.clone(),
         rpar: generator_exp.rpar.clone(),
@@ -101,8 +106,13 @@ pub fn fix_unnecessary_generator_set(locator: &SourceCodeLocator, expr: &Expr) -
     } else {
         return Err(anyhow::anyhow!("Expected node to be: Expression::Call."));
     };
-    let arg = if let Some(Arg { value, .. }) = call.args.first_mut() {
+    let (arg, whitespace_after_arg) = if let Some(Arg {
-        value
+        value,
+        whitespace_after_arg,
+        ..
+    }) = call.args.first_mut()
+    {
+        (value, whitespace_after_arg)
     } else {
         return Err(anyhow::anyhow!("Expected node to be: Arg."));
     };
@@ -121,7 +131,7 @@ pub fn fix_unnecessary_generator_set(locator: &SourceCodeLocator, expr: &Expr) -
             whitespace_after: call.whitespace_before_args.clone(),
         },
         rbrace: RightCurlyBrace {
-            whitespace_before: call.whitespace_after_func.clone(),
+            whitespace_before: whitespace_after_arg.clone(),
         },
         lpar: generator_exp.lpar.clone(),
         rpar: generator_exp.rpar.clone(),
@@ -168,8 +178,13 @@ pub fn fix_unnecessary_list_comprehension_set(
     } else {
         return Err(anyhow::anyhow!("Expected node to be: Expression::Call."));
     };
-    let arg = if let Some(Arg { value, .. }) = call.args.first_mut() {
+    let (arg, whitespace_after_arg) = if let Some(Arg {
-        value
+        value,
+        whitespace_after_arg,
+        ..
+    }) = call.args.first_mut()
+    {
+        (value, whitespace_after_arg)
     } else {
         return Err(anyhow::anyhow!("Expected node to be: Arg."));
     };
@@ -188,7 +203,7 @@ pub fn fix_unnecessary_list_comprehension_set(
             whitespace_after: call.whitespace_before_args.clone(),
         },
         rbrace: RightCurlyBrace {
-            whitespace_before: call.whitespace_after_func.clone(),
+            whitespace_before: whitespace_after_arg.clone(),
         },
         lpar: list_comp.lpar.clone(),
         rpar: list_comp.rpar.clone(),
@@ -232,8 +247,13 @@ pub fn fix_unnecessary_literal_set(locator: &SourceCodeLocator, expr: &Expr) ->
     } else {
         return Err(anyhow::anyhow!("Expected node to be: Expression::Call."));
     };
-    let arg = if let Some(Arg { value, .. }) = call.args.first_mut() {
+    let (arg, whitespace_after_arg) = if let Some(Arg {
-        value
+        value,
+        whitespace_after_arg,
+        ..
+    }) = call.args.first_mut()
+    {
+        (value, whitespace_after_arg)
     } else {
         return Err(anyhow::anyhow!("Expected node to be: Arg."));
     };
@@ -256,7 +276,7 @@ pub fn fix_unnecessary_literal_set(locator: &SourceCodeLocator, expr: &Expr) ->
                 whitespace_after: call.whitespace_before_args.clone(),
             },
             rbrace: RightCurlyBrace {
-                whitespace_before: call.whitespace_after_func.clone(),
+                whitespace_before: whitespace_after_arg.clone(),
             },
             lpar: Default::default(),
             rpar: Default::default(),

src/snapshots/ruff__linter__tests__C400_C400.py.snap

@@ -19,4 +19,21 @@ expression: checks
         row: 1
         column: 29
     applied: false
+- kind: UnnecessaryGeneratorList
+  location:
+    row: 2
+    column: 4
+  end_location:
+    row: 4
+    column: 1
+  fix:
+    patch:
+      content: "[\n    x for x in range(3)\n]"
+      location:
+        row: 2
+        column: 4
+      end_location:
+        row: 4
+        column: 1
+    applied: false
 

src/snapshots/ruff__linter__tests__C401_C401.py.snap

@@ -19,4 +19,21 @@ expression: checks
         row: 1
         column: 28
     applied: false
+- kind: UnnecessaryGeneratorSet
+  location:
+    row: 2
+    column: 4
+  end_location:
+    row: 4
+    column: 1
+  fix:
+    patch:
+      content: "{\n    x for x in range(3)\n}"
+      location:
+        row: 2
+        column: 4
+      end_location:
+        row: 4
+        column: 1
+    applied: false
 

src/snapshots/ruff__linter__tests__C403_C403.py.snap

@@ -19,4 +19,21 @@ expression: checks
         row: 1
         column: 30
     applied: false
+- kind: UnnecessaryListComprehensionSet
+  location:
+    row: 2
+    column: 4
+  end_location:
+    row: 4
+    column: 1
+  fix:
+    patch:
+      content: "{\n    x for x in range(3)\n}"
+      location:
+        row: 2
+        column: 4
+      end_location:
+        row: 4
+        column: 1
+    applied: false