diff --git a/README.md b/README.md index 9aa129f61e..a013ee35cd 100644 --- a/README.md +++ b/README.md @@ -453,7 +453,7 @@ The 🛠 emoji indicates that a rule is automatically fixable by the `--fix` com | C413 | UnnecessaryCallAroundSorted | Unnecessary `(list\|reversed)` call around `sorted()` | | | C414 | UnnecessaryDoubleCastOrProcess | Unnecessary `(list\|reversed\|set\|sorted\|tuple)` call within `(list\|set\|sorted\|tuple)()` | | | C415 | UnnecessarySubscriptReversal | Unnecessary subscript reversal of iterable within `(reversed\|set\|sorted)()` | | -| C416 | UnnecessaryComprehension | Unnecessary `(list\|set)` comprehension (rewrite using `(list\|set)()`) | | +| C416 | UnnecessaryComprehension | Unnecessary `(list\|set)` comprehension (rewrite using `(list\|set)()`) | 🛠 | | C417 | UnnecessaryMap | Unnecessary `map` usage (rewrite using a `(list\|set\|dict)` comprehension) | | ### flake8-bugbear diff --git a/src/check_ast.rs b/src/check_ast.rs index 9ed8f3d8b2..aa48ad756c 100644 --- a/src/check_ast.rs +++ b/src/check_ast.rs @@ -1363,6 +1363,8 @@ where expr, elt, generators, + self.locator, + self.patch(), self.locate_check(Range::from_located(expr)), ) { self.checks.push(check); diff --git a/src/checks.rs b/src/checks.rs index 53420a0f92..5afd0bd9bb 100644 --- a/src/checks.rs +++ b/src/checks.rs @@ -1470,6 +1470,7 @@ impl CheckKind { | CheckKind::TypeOfPrimitive(_) | CheckKind::UnnecessaryAbspath | CheckKind::UnnecessaryCollectionCall(_) + | CheckKind::UnnecessaryComprehension(_) | CheckKind::UnnecessaryGeneratorDict | CheckKind::UnnecessaryGeneratorList | CheckKind::UnnecessaryGeneratorSet diff --git a/src/flake8_comprehensions/checks.rs b/src/flake8_comprehensions/checks.rs index b6ede045c0..54b3891f57 100644 --- a/src/flake8_comprehensions/checks.rs +++ b/src/flake8_comprehensions/checks.rs @@ -443,6 +443,8 @@ pub fn unnecessary_comprehension( expr: &Expr, elt: &Expr, generators: &[Comprehension], + locator: &SourceCodeLocator, + fix: bool, location: Range, ) -> Option { if generators.len() != 1 { @@ -462,10 +464,17 @@ pub fn unnecessary_comprehension( ExprKind::SetComp { .. } => "set", _ => return None, }; - Some(Check::new( + let mut check = Check::new( CheckKind::UnnecessaryComprehension(expr_kind.to_string()), location, - )) + ); + if fix { + match fixes::fix_unnecessary_comprehension(locator, expr) { + Ok(fix) => check.amend(fix), + Err(e) => error!("Failed to generate fix: {}", e), + } + } + Some(check) } /// C417 diff --git a/src/flake8_comprehensions/fixes.rs b/src/flake8_comprehensions/fixes.rs index c5988f2a73..34041ec56c 100644 --- a/src/flake8_comprehensions/fixes.rs +++ b/src/flake8_comprehensions/fixes.rs @@ -1,7 +1,7 @@ use anyhow::Result; use libcst_native::{ Arg, Call, Codegen, Dict, DictComp, DictElement, Element, Expr, Expression, LeftCurlyBrace, - LeftParen, LeftSquareBracket, List, ListComp, Module, ParenthesizableWhitespace, + LeftParen, LeftSquareBracket, List, ListComp, Module, Name, ParenthesizableWhitespace, RightCurlyBrace, RightParen, RightSquareBracket, Set, SetComp, SimpleString, SimpleWhitespace, SmallStatement, Statement, Tuple, }; @@ -503,7 +503,7 @@ pub fn fix_unnecessary_literal_within_list_call( )) } -/// (C411) Convert `list([i for i in x])` to `[i for i in x]`. +/// (C411) Convert `list([i * i for i in x])` to `[i * i for i in x]`. pub fn fix_unnecessary_list_call( locator: &SourceCodeLocator, expr: &rustpython_ast::Expr, @@ -525,3 +525,73 @@ pub fn fix_unnecessary_list_call( expr.end_location.unwrap(), )) } + +/// (C416) Convert `[i for i in x]` to `list(x)`. +pub fn fix_unnecessary_comprehension( + locator: &SourceCodeLocator, + expr: &rustpython_ast::Expr, +) -> Result { + let mut tree = match_tree(locator, expr)?; + let mut body = match_expr(&mut tree)?; + + match &body.value { + Expression::ListComp(inner) => { + body.value = Expression::Call(Box::new(Call { + func: Box::new(Expression::Name(Box::new(Name { + value: "list", + lpar: Default::default(), + rpar: Default::default(), + }))), + args: vec![Arg { + value: inner.for_in.iter.clone(), + keyword: Default::default(), + equal: Default::default(), + comma: Default::default(), + star: Default::default(), + whitespace_after_star: Default::default(), + whitespace_after_arg: Default::default(), + }], + lpar: Default::default(), + rpar: Default::default(), + whitespace_after_func: Default::default(), + whitespace_before_args: Default::default(), + })) + } + Expression::SetComp(inner) => { + body.value = Expression::Call(Box::new(Call { + func: Box::new(Expression::Name(Box::new(Name { + value: "set", + lpar: Default::default(), + rpar: Default::default(), + }))), + args: vec![Arg { + value: inner.for_in.iter.clone(), + keyword: Default::default(), + equal: Default::default(), + comma: Default::default(), + star: Default::default(), + whitespace_after_star: Default::default(), + whitespace_after_arg: Default::default(), + }], + lpar: Default::default(), + rpar: Default::default(), + whitespace_after_func: Default::default(), + whitespace_before_args: Default::default(), + })) + } + _ => { + return Err(anyhow::anyhow!( + "Expected node to be: Expression::ListComp | Expression:SetComp." + )) + } + } + + let mut state = Default::default(); + tree.codegen(&mut state); + + Ok(Fix::replacement( + state.to_string(), + expr.location, + expr.end_location.unwrap(), + )) +} diff --git a/src/snapshots/ruff__linter__tests__C416_C416.py.snap b/src/snapshots/ruff__linter__tests__C416_C416.py.snap index d5ee71b449..12866d5074 100644 --- a/src/snapshots/ruff__linter__tests__C416_C416.py.snap +++ b/src/snapshots/ruff__linter__tests__C416_C416.py.snap @@ -10,7 +10,16 @@ expression: checks end_location: row: 2 column: 14 - fix: ~ + fix: + patch: + content: list(x) + location: + row: 2 + column: 0 + end_location: + row: 2 + column: 14 + applied: false - kind: UnnecessaryComprehension: set location: @@ -19,5 +28,14 @@ expression: checks end_location: row: 3 column: 14 - fix: ~ + fix: + patch: + content: set(x) + location: + row: 3 + column: 0 + end_location: + row: 3 + column: 14 + applied: false