diff --git a/resources/test/fixtures/C400.py b/resources/test/fixtures/C400.py index 419203b4fe..6f19bd197a 100644 --- a/resources/test/fixtures/C400.py +++ b/resources/test/fixtures/C400.py @@ -1 +1,4 @@ x = list(x for x in range(3)) +x = list( + x for x in range(3) +) diff --git a/resources/test/fixtures/C401.py b/resources/test/fixtures/C401.py index b6b085d711..3842429f9e 100644 --- a/resources/test/fixtures/C401.py +++ b/resources/test/fixtures/C401.py @@ -1 +1,4 @@ x = set(x for x in range(3)) +x = set( + x for x in range(3) +) diff --git a/resources/test/fixtures/C403.py b/resources/test/fixtures/C403.py index 778cad32ca..51e8179d8c 100644 --- a/resources/test/fixtures/C403.py +++ b/resources/test/fixtures/C403.py @@ -1 +1,4 @@ s = set([x for x in range(3)]) +s = set( + [x for x in range(3)] +) diff --git a/src/flake8_comprehensions/fixes.rs b/src/flake8_comprehensions/fixes.rs index 9b4e21232c..075aa00a01 100644 --- a/src/flake8_comprehensions/fixes.rs +++ b/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() { - value + let (arg, whitespace_after_arg) = if let Some(Arg { + 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() { - value + let (arg, whitespace_after_arg) = if let Some(Arg { + 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() { - value + let (arg, whitespace_after_arg) = if let Some(Arg { + 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() { - value + let (arg, whitespace_after_arg) = if let Some(Arg { + 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(), diff --git a/src/snapshots/ruff__linter__tests__C400_C400.py.snap b/src/snapshots/ruff__linter__tests__C400_C400.py.snap index c7cec50ee1..638eeebfea 100644 --- a/src/snapshots/ruff__linter__tests__C400_C400.py.snap +++ b/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 diff --git a/src/snapshots/ruff__linter__tests__C401_C401.py.snap b/src/snapshots/ruff__linter__tests__C401_C401.py.snap index 58d0cb0c51..63419cbb84 100644 --- a/src/snapshots/ruff__linter__tests__C401_C401.py.snap +++ b/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 diff --git a/src/snapshots/ruff__linter__tests__C403_C403.py.snap b/src/snapshots/ruff__linter__tests__C403_C403.py.snap index a71b386c44..c152e4d6c5 100644 --- a/src/snapshots/ruff__linter__tests__C403_C403.py.snap +++ b/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