diff --git a/crates/ruff/src/autofix/actions.rs b/crates/ruff/src/autofix/actions.rs index 283726cede..afa61a274f 100644 --- a/crates/ruff/src/autofix/actions.rs +++ b/crates/ruff/src/autofix/actions.rs @@ -14,7 +14,7 @@ use ruff_python_ast::newlines::NewlineWithTrailingNewline; use ruff_python_ast::source_code::{Indexer, Locator, Stylist}; use crate::cst::helpers::compose_module_path; -use crate::cst::matchers::match_module; +use crate::cst::matchers::match_statement; /// Determine if a body contains only a single statement, taking into account /// deleted. @@ -212,9 +212,9 @@ pub(crate) fn remove_unused_imports<'a>( stylist: &Stylist, ) -> Result { let module_text = locator.slice(stmt.range()); - let mut tree = match_module(module_text)?; + let mut tree = match_statement(module_text)?; - let Some(Statement::Simple(body)) = tree.body.first_mut() else { + let Statement::Simple(body) = &mut tree else { bail!("Expected Statement::Simple"); }; diff --git a/crates/ruff/src/cst/matchers.rs b/crates/ruff/src/cst/matchers.rs index d24c0580d8..0df0590fcc 100644 --- a/crates/ruff/src/cst/matchers.rs +++ b/crates/ruff/src/cst/matchers.rs @@ -1,8 +1,9 @@ use anyhow::{bail, Result}; use libcst_native::{ - Attribute, Call, Comparison, Dict, Expr, Expression, FormattedString, FormattedStringContent, - FormattedStringExpression, Import, ImportAlias, ImportFrom, ImportNames, Module, Name, - SimpleString, SmallStatement, Statement, + Arg, Attribute, Call, Comparison, CompoundStatement, Dict, Expression, FormattedString, + FormattedStringContent, FormattedStringExpression, FunctionDef, GeneratorExp, If, Import, + ImportAlias, ImportFrom, ImportNames, IndentedBlock, Lambda, ListComp, Module, Name, + SimpleString, SmallStatement, Statement, Suite, Tuple, With, }; pub(crate) fn match_module(module_text: &str) -> Result { @@ -19,20 +20,15 @@ pub(crate) fn match_expression(expression_text: &str) -> Result { } } -pub(crate) fn match_expr<'a, 'b>(module: &'a mut Module<'b>) -> Result<&'a mut Expr<'b>> { - if let Some(Statement::Simple(expr)) = module.body.first_mut() { - if let Some(SmallStatement::Expr(expr)) = expr.body.first_mut() { - Ok(expr) - } else { - bail!("Expected SmallStatement::Expr") - } - } else { - bail!("Expected Statement::Simple") +pub(crate) fn match_statement(statement_text: &str) -> Result { + match libcst_native::parse_statement(statement_text) { + Ok(statement) => Ok(statement), + Err(_) => bail!("Failed to extract statement from source"), } } -pub(crate) fn match_import<'a, 'b>(module: &'a mut Module<'b>) -> Result<&'a mut Import<'b>> { - if let Some(Statement::Simple(expr)) = module.body.first_mut() { +pub(crate) fn match_import<'a, 'b>(statement: &'a mut Statement<'b>) -> Result<&'a mut Import<'b>> { + if let Statement::Simple(expr) = statement { if let Some(SmallStatement::Import(expr)) = expr.body.first_mut() { Ok(expr) } else { @@ -44,9 +40,9 @@ pub(crate) fn match_import<'a, 'b>(module: &'a mut Module<'b>) -> Result<&'a mut } pub(crate) fn match_import_from<'a, 'b>( - module: &'a mut Module<'b>, + statement: &'a mut Statement<'b>, ) -> Result<&'a mut ImportFrom<'b>> { - if let Some(Statement::Simple(expr)) = module.body.first_mut() { + if let Statement::Simple(expr) = statement { if let Some(SmallStatement::ImportFrom(expr)) = expr.body.first_mut() { Ok(expr) } else { @@ -67,7 +63,17 @@ pub(crate) fn match_aliases<'a, 'b>( } } -pub(crate) fn match_call<'a, 'b>(expression: &'a mut Expression<'b>) -> Result<&'a mut Call<'b>> { +pub(crate) fn match_call<'a, 'b>(expression: &'a Expression<'b>) -> Result<&'a Call<'b>> { + if let Expression::Call(call) = expression { + Ok(call) + } else { + bail!("Expected Expression::Call") + } +} + +pub(crate) fn match_call_mut<'a, 'b>( + expression: &'a mut Expression<'b>, +) -> Result<&'a mut Call<'b>> { if let Expression::Call(call) = expression { Ok(call) } else { @@ -135,10 +141,100 @@ pub(crate) fn match_formatted_string_expression<'a, 'b>( } } -pub(crate) fn match_name<'a, 'b>(expression: &'a mut Expression<'b>) -> Result<&'a mut Name<'b>> { +pub(crate) fn match_name<'a, 'b>(expression: &'a Expression<'b>) -> Result<&'a Name<'b>> { if let Expression::Name(name) = expression { Ok(name) } else { bail!("Expected Expression::Name") } } + +pub(crate) fn match_arg<'a, 'b>(call: &'a Call<'b>) -> Result<&'a Arg<'b>> { + if let Some(arg) = call.args.first() { + Ok(arg) + } else { + bail!("Expected Arg") + } +} + +pub(crate) fn match_generator_exp<'a, 'b>( + expression: &'a Expression<'b>, +) -> Result<&'a GeneratorExp<'b>> { + if let Expression::GeneratorExp(generator_exp) = expression { + Ok(generator_exp) + } else { + bail!("Expected Expression::GeneratorExp") + } +} + +pub(crate) fn match_tuple<'a, 'b>(expression: &'a Expression<'b>) -> Result<&'a Tuple<'b>> { + if let Expression::Tuple(tuple) = expression { + Ok(tuple) + } else { + bail!("Expected Expression::Tuple") + } +} + +pub(crate) fn match_list_comp<'a, 'b>(expression: &'a Expression<'b>) -> Result<&'a ListComp<'b>> { + if let Expression::ListComp(list_comp) = expression { + Ok(list_comp) + } else { + bail!("Expected Expression::ListComp") + } +} + +pub(crate) fn match_lambda<'a, 'b>(expression: &'a Expression<'b>) -> Result<&'a Lambda<'b>> { + if let Expression::Lambda(lambda) = expression { + Ok(lambda) + } else { + bail!("Expected Expression::Lambda") + } +} + +pub(crate) fn match_function_def<'a, 'b>( + statement: &'a mut Statement<'b>, +) -> Result<&'a mut FunctionDef<'b>> { + if let Statement::Compound(compound) = statement { + if let CompoundStatement::FunctionDef(function_def) = compound { + Ok(function_def) + } else { + bail!("Expected CompoundStatement::FunctionDef") + } + } else { + bail!("Expected Statement::Compound") + } +} + +pub(crate) fn match_indented_block<'a, 'b>( + suite: &'a mut Suite<'b>, +) -> Result<&'a mut IndentedBlock<'b>> { + if let Suite::IndentedBlock(indented_block) = suite { + Ok(indented_block) + } else { + bail!("Expected Suite::IndentedBlock") + } +} + +pub(crate) fn match_with<'a, 'b>(statement: &'a mut Statement<'b>) -> Result<&'a mut With<'b>> { + if let Statement::Compound(compound) = statement { + if let CompoundStatement::With(with) = compound { + Ok(with) + } else { + bail!("Expected CompoundStatement::With") + } + } else { + bail!("Expected Statement::Compound") + } +} + +pub(crate) fn match_if<'a, 'b>(statement: &'a mut Statement<'b>) -> Result<&'a mut If<'b>> { + if let Statement::Compound(compound) = statement { + if let CompoundStatement::If(if_) = compound { + Ok(if_) + } else { + bail!("Expected CompoundStatement::If") + } + } else { + bail!("Expected Statement::Compound") + } +} diff --git a/crates/ruff/src/importer/mod.rs b/crates/ruff/src/importer/mod.rs index a694bd6f74..7179459c5b 100644 --- a/crates/ruff/src/importer/mod.rs +++ b/crates/ruff/src/importer/mod.rs @@ -10,7 +10,7 @@ use ruff_python_ast::imports::{AnyImport, Import}; use ruff_python_ast::source_code::{Locator, Stylist}; use ruff_python_semantic::model::SemanticModel; -use crate::cst::matchers::{match_aliases, match_import_from, match_module}; +use crate::cst::matchers::{match_aliases, match_import_from, match_statement}; use crate::importer::insertion::Insertion; mod insertion; @@ -199,8 +199,8 @@ impl<'a> Importer<'a> { /// Add the given member to an existing `Stmt::ImportFrom` statement. fn add_member(&self, stmt: &Stmt, member: &str) -> Result { - let mut tree = match_module(self.locator.slice(stmt.range()))?; - let import_from = match_import_from(&mut tree)?; + let mut statement = match_statement(self.locator.slice(stmt.range()))?; + let import_from = match_import_from(&mut statement)?; let aliases = match_aliases(import_from)?; aliases.push(ImportAlias { name: NameOrAttribute::N(Box::new(Name { @@ -216,7 +216,7 @@ impl<'a> Importer<'a> { default_indent: self.stylist.indentation(), ..CodegenState::default() }; - tree.codegen(&mut state); + statement.codegen(&mut state); Ok(Edit::range_replacement(state.to_string(), stmt.range())) } } diff --git a/crates/ruff/src/rules/flake8_comprehensions/fixes.rs b/crates/ruff/src/rules/flake8_comprehensions/fixes.rs index 24deb88e1e..b9054be763 100644 --- a/crates/ruff/src/rules/flake8_comprehensions/fixes.rs +++ b/crates/ruff/src/rules/flake8_comprehensions/fixes.rs @@ -2,33 +2,20 @@ use anyhow::{bail, Result}; use itertools::Itertools; use libcst_native::{ Arg, AssignEqual, AssignTargetExpression, Call, Codegen, CodegenState, Comment, CompFor, Dict, - DictComp, DictElement, Element, EmptyLine, Expr, Expression, GeneratorExp, LeftCurlyBrace, - LeftParen, LeftSquareBracket, List, ListComp, Name, ParenthesizableWhitespace, - ParenthesizedWhitespace, RightCurlyBrace, RightParen, RightSquareBracket, Set, SetComp, - SimpleString, SimpleWhitespace, TrailingWhitespace, Tuple, + DictComp, DictElement, Element, EmptyLine, Expression, GeneratorExp, LeftCurlyBrace, LeftParen, + LeftSquareBracket, List, ListComp, Name, ParenthesizableWhitespace, ParenthesizedWhitespace, + RightCurlyBrace, RightParen, RightSquareBracket, Set, SetComp, SimpleString, SimpleWhitespace, + TrailingWhitespace, Tuple, }; use rustpython_parser::ast::Ranged; use ruff_diagnostics::{Edit, Fix}; use ruff_python_ast::source_code::{Locator, Stylist}; -use crate::cst::matchers::{match_expr, match_module}; - -fn match_call<'a, 'b>(expr: &'a mut Expr<'b>) -> Result<&'a mut Call<'b>> { - if let Expression::Call(call) = &mut expr.value { - Ok(call) - } else { - bail!("Expected Expression::Call") - } -} - -fn match_arg<'a, 'b>(call: &'a Call<'b>) -> Result<&'a Arg<'b>> { - if let Some(arg) = call.args.first() { - Ok(arg) - } else { - bail!("Expected Arg") - } -} +use crate::cst::matchers::{ + match_arg, match_call, match_call_mut, match_expression, match_generator_exp, match_lambda, + match_list_comp, match_name, match_tuple, +}; /// (C400) Convert `list(x for x in y)` to `[x for x in y]`. pub(crate) fn fix_unnecessary_generator_list( @@ -38,18 +25,13 @@ pub(crate) fn fix_unnecessary_generator_list( ) -> Result { // Expr(Call(GeneratorExp)))) -> Expr(ListComp))) let module_text = locator.slice(expr.range()); - let mut tree = match_module(module_text)?; - let mut body = match_expr(&mut tree)?; - let call = match_call(body)?; + let mut tree = match_expression(module_text)?; + let call = match_call_mut(&mut tree)?; let arg = match_arg(call)?; - let Expression::GeneratorExp(generator_exp) = &arg.value else { - bail!( - "Expected Expression::GeneratorExp" - ); - }; + let generator_exp = match_generator_exp(&arg.value)?; - body.value = Expression::ListComp(Box::new(ListComp { + tree = Expression::ListComp(Box::new(ListComp { elt: generator_exp.elt.clone(), for_in: generator_exp.for_in.clone(), lbracket: LeftSquareBracket { @@ -81,18 +63,13 @@ pub(crate) fn fix_unnecessary_generator_set( ) -> Result { // Expr(Call(GeneratorExp)))) -> Expr(SetComp))) let module_text = locator.slice(expr.range()); - let mut tree = match_module(module_text)?; - let mut body = match_expr(&mut tree)?; - let call = match_call(body)?; + let mut tree = match_expression(module_text)?; + let call = match_call_mut(&mut tree)?; let arg = match_arg(call)?; - let Expression::GeneratorExp(generator_exp) = &arg.value else { - bail!( - "Expected Expression::GeneratorExp" - ); - }; + let generator_exp = match_generator_exp(&arg.value)?; - body.value = Expression::SetComp(Box::new(SetComp { + tree = Expression::SetComp(Box::new(SetComp { elt: generator_exp.elt.clone(), for_in: generator_exp.for_in.clone(), lbrace: LeftCurlyBrace { @@ -132,32 +109,18 @@ pub(crate) fn fix_unnecessary_generator_dict( parent: Option<&rustpython_parser::ast::Expr>, ) -> Result { let module_text = locator.slice(expr.range()); - let mut tree = match_module(module_text)?; - let mut body = match_expr(&mut tree)?; - let call = match_call(body)?; + let mut tree = match_expression(module_text)?; + let call = match_call_mut(&mut tree)?; let arg = match_arg(call)?; // Extract the (k, v) from `(k, v) for ...`. - let Expression::GeneratorExp(generator_exp) = &arg.value else { - bail!( - "Expected Expression::GeneratorExp" - ); - }; - let Expression::Tuple(tuple) = &generator_exp.elt.as_ref() else { - bail!("Expected Expression::Tuple"); - }; - let Some(Element::Simple { value: key, .. }) = &tuple.elements.get(0) else { - bail!( - "Expected tuple to contain a key as the first element" - ); - }; - let Some(Element::Simple { value, .. }) = &tuple.elements.get(1) else { - bail!( - "Expected tuple to contain a key as the second element" - ); + let generator_exp = match_generator_exp(&arg.value)?; + let tuple = match_tuple(&generator_exp.elt)?; + let [Element::Simple { value: key, .. }, Element::Simple { value, .. }] = &tuple.elements[..] else { + bail!("Expected tuple to contain two elements"); }; - body.value = Expression::DictComp(Box::new(DictComp { + tree = Expression::DictComp(Box::new(DictComp { key: Box::new(key.clone()), value: Box::new(value.clone()), for_in: generator_exp.for_in.clone(), @@ -200,16 +163,13 @@ pub(crate) fn fix_unnecessary_list_comprehension_set( // Expr(Call(ListComp)))) -> // Expr(SetComp))) let module_text = locator.slice(expr.range()); - let mut tree = match_module(module_text)?; - let mut body = match_expr(&mut tree)?; - let call = match_call(body)?; + let mut tree = match_expression(module_text)?; + let call = match_call_mut(&mut tree)?; let arg = match_arg(call)?; - let Expression::ListComp(list_comp) = &arg.value else { - bail!("Expected Expression::ListComp"); - }; + let list_comp = match_list_comp(&arg.value)?; - body.value = Expression::SetComp(Box::new(SetComp { + tree = Expression::SetComp(Box::new(SetComp { elt: list_comp.elt.clone(), for_in: list_comp.for_in.clone(), lbrace: LeftCurlyBrace { @@ -240,25 +200,20 @@ pub(crate) fn fix_unnecessary_list_comprehension_dict( expr: &rustpython_parser::ast::Expr, ) -> Result { let module_text = locator.slice(expr.range()); - let mut tree = match_module(module_text)?; - let mut body = match_expr(&mut tree)?; - let call = match_call(body)?; + let mut tree = match_expression(module_text)?; + let call = match_call_mut(&mut tree)?; let arg = match_arg(call)?; - let Expression::ListComp(list_comp) = &arg.value else { - bail!("Expected Expression::ListComp") - }; + let list_comp = match_list_comp(&arg.value)?; - let Expression::Tuple(tuple) = &*list_comp.elt else { - bail!("Expected Expression::Tuple") - }; + let tuple = match_tuple(&list_comp.elt)?; let [Element::Simple { value: key, comma: Some(comma), }, Element::Simple { value, .. }] = &tuple.elements[..] else { bail!("Expected tuple with two elements"); }; - body.value = Expression::DictComp(Box::new(DictComp { + tree = Expression::DictComp(Box::new(DictComp { key: Box::new(key.clone()), value: Box::new(value.clone()), for_in: list_comp.for_in.clone(), @@ -335,9 +290,8 @@ pub(crate) fn fix_unnecessary_literal_set( ) -> Result { // Expr(Call(List|Tuple)))) -> Expr(Set))) let module_text = locator.slice(expr.range()); - let mut tree = match_module(module_text)?; - let mut body = match_expr(&mut tree)?; - let mut call = match_call(body)?; + let mut tree = match_expression(module_text)?; + let mut call = match_call_mut(&mut tree)?; let arg = match_arg(call)?; let (elements, whitespace_after, whitespace_before) = match &arg.value { @@ -355,7 +309,7 @@ pub(crate) fn fix_unnecessary_literal_set( if elements.is_empty() { call.args = vec![]; } else { - body.value = Expression::Set(Box::new(Set { + tree = Expression::Set(Box::new(Set { elements, lbrace: LeftCurlyBrace { whitespace_after }, rbrace: RightCurlyBrace { whitespace_before }, @@ -382,9 +336,8 @@ pub(crate) fn fix_unnecessary_literal_dict( ) -> Result { // Expr(Call(List|Tuple)))) -> Expr(Dict))) let module_text = locator.slice(expr.range()); - let mut tree = match_module(module_text)?; - let mut body = match_expr(&mut tree)?; - let call = match_call(body)?; + let mut tree = match_expression(module_text)?; + let call = match_call_mut(&mut tree)?; let arg = match_arg(call)?; let elements = match &arg.value { @@ -421,7 +374,7 @@ pub(crate) fn fix_unnecessary_literal_dict( }) .collect::>>()?; - body.value = Expression::Dict(Box::new(Dict { + tree = Expression::Dict(Box::new(Dict { elements, lbrace: LeftCurlyBrace { whitespace_after: call.whitespace_before_args.clone(), @@ -451,12 +404,9 @@ pub(crate) fn fix_unnecessary_collection_call( ) -> Result { // Expr(Call("list" | "tuple" | "dict")))) -> Expr(List|Tuple|Dict) let module_text = locator.slice(expr.range()); - let mut tree = match_module(module_text)?; - let mut body = match_expr(&mut tree)?; - let call = match_call(body)?; - let Expression::Name(name) = &call.func.as_ref() else { - bail!("Expected Expression::Name"); - }; + let mut tree = match_expression(module_text)?; + let call = match_call_mut(&mut tree)?; + let name = match_name(&call.func)?; // Arena allocator used to create formatted strings of sufficient lifetime, // below. @@ -464,14 +414,14 @@ pub(crate) fn fix_unnecessary_collection_call( match name.value { "tuple" => { - body.value = Expression::Tuple(Box::new(Tuple { + tree = Expression::Tuple(Box::new(Tuple { elements: vec![], lpar: vec![LeftParen::default()], rpar: vec![RightParen::default()], })); } "list" => { - body.value = Expression::List(Box::new(List { + tree = Expression::List(Box::new(List { elements: vec![], lbracket: LeftSquareBracket::default(), rbracket: RightSquareBracket::default(), @@ -481,7 +431,7 @@ pub(crate) fn fix_unnecessary_collection_call( } "dict" => { if call.args.is_empty() { - body.value = Expression::Dict(Box::new(Dict { + tree = Expression::Dict(Box::new(Dict { elements: vec![], lbrace: LeftCurlyBrace::default(), rbrace: RightCurlyBrace::default(), @@ -522,7 +472,7 @@ pub(crate) fn fix_unnecessary_collection_call( }) .collect(); - body.value = Expression::Dict(Box::new(Dict { + tree = Expression::Dict(Box::new(Dict { elements, lbrace: LeftCurlyBrace { whitespace_after: call.whitespace_before_args.clone(), @@ -562,9 +512,8 @@ pub(crate) fn fix_unnecessary_literal_within_tuple_call( expr: &rustpython_parser::ast::Expr, ) -> Result { let module_text = locator.slice(expr.range()); - let mut tree = match_module(module_text)?; - let mut body = match_expr(&mut tree)?; - let call = match_call(body)?; + let mut tree = match_expression(module_text)?; + let call = match_call_mut(&mut tree)?; let arg = match_arg(call)?; let (elements, whitespace_after, whitespace_before) = match &arg.value { Expression::Tuple(inner) => ( @@ -590,7 +539,7 @@ pub(crate) fn fix_unnecessary_literal_within_tuple_call( } }; - body.value = Expression::Tuple(Box::new(Tuple { + tree = Expression::Tuple(Box::new(Tuple { elements: elements.clone(), lpar: vec![LeftParen { whitespace_after: whitespace_after.clone(), @@ -617,9 +566,8 @@ pub(crate) fn fix_unnecessary_literal_within_list_call( expr: &rustpython_parser::ast::Expr, ) -> Result { let module_text = locator.slice(expr.range()); - let mut tree = match_module(module_text)?; - let mut body = match_expr(&mut tree)?; - let call = match_call(body)?; + let mut tree = match_expression(module_text)?; + let call = match_call_mut(&mut tree)?; let arg = match_arg(call)?; let (elements, whitespace_after, whitespace_before) = match &arg.value { Expression::Tuple(inner) => ( @@ -645,7 +593,7 @@ pub(crate) fn fix_unnecessary_literal_within_list_call( } }; - body.value = Expression::List(Box::new(List { + tree = Expression::List(Box::new(List { elements: elements.clone(), lbracket: LeftSquareBracket { whitespace_after: whitespace_after.clone(), @@ -675,12 +623,11 @@ pub(crate) fn fix_unnecessary_list_call( ) -> Result { // Expr(Call(List|Tuple)))) -> Expr(List|Tuple))) let module_text = locator.slice(expr.range()); - let mut tree = match_module(module_text)?; - let mut body = match_expr(&mut tree)?; - let call = match_call(body)?; + let mut tree = match_expression(module_text)?; + let call = match_call_mut(&mut tree)?; let arg = match_arg(call)?; - body.value = arg.value.clone(); + tree = arg.value.clone(); let mut state = CodegenState { default_newline: &stylist.line_ending(), @@ -701,17 +648,10 @@ pub(crate) fn fix_unnecessary_call_around_sorted( expr: &rustpython_parser::ast::Expr, ) -> Result { let module_text = locator.slice(expr.range()); - let mut tree = match_module(module_text)?; - let mut body = match_expr(&mut tree)?; - let outer_call = match_call(body)?; + let mut tree = match_expression(module_text)?; + let outer_call = match_call_mut(&mut tree)?; let inner_call = match &outer_call.args[..] { - [arg] => { - if let Expression::Call(call) = &arg.value { - call - } else { - bail!("Expected Expression::Call "); - } - } + [arg] => match_call(&arg.value)?, _ => { bail!("Expected one argument in outer function call"); } @@ -719,7 +659,7 @@ pub(crate) fn fix_unnecessary_call_around_sorted( if let Expression::Name(outer_name) = &*outer_call.func { if outer_name.value == "list" { - body.value = Expression::Call(inner_call.clone()); + tree = Expression::Call(Box::new((*inner_call).clone())); } else { // If the `reverse` argument is used let args = if inner_call.args.iter().any(|arg| { @@ -796,7 +736,7 @@ pub(crate) fn fix_unnecessary_call_around_sorted( args }; - body.value = Expression::Call(Box::new(Call { + tree = Expression::Call(Box::new(Call { func: inner_call.func.clone(), args, lpar: inner_call.lpar.clone(), @@ -824,15 +764,12 @@ pub(crate) fn fix_unnecessary_double_cast_or_process( expr: &rustpython_parser::ast::Expr, ) -> Result { let module_text = locator.slice(expr.range()); - let mut tree = match_module(module_text)?; - let body = match_expr(&mut tree)?; - let mut outer_call = match_call(body)?; + let mut tree = match_expression(module_text)?; + let mut outer_call = match_call_mut(&mut tree)?; outer_call.args = match outer_call.args.split_first() { Some((first, rest)) => { - let Expression::Call(inner_call) = &first.value else { - bail!("Expected Expression::Call "); - }; + let inner_call = match_call(&first.value)?; if let Some(iterable) = inner_call.args.first() { let mut args = vec![iterable.clone()]; args.extend_from_slice(rest); @@ -861,12 +798,11 @@ pub(crate) fn fix_unnecessary_comprehension( expr: &rustpython_parser::ast::Expr, ) -> Result { let module_text = locator.slice(expr.range()); - let mut tree = match_module(module_text)?; - let mut body = match_expr(&mut tree)?; + let mut tree = match_expression(module_text)?; - match &body.value { + match &tree { Expression::ListComp(inner) => { - body.value = Expression::Call(Box::new(Call { + tree = Expression::Call(Box::new(Call { func: Box::new(Expression::Name(Box::new(Name { value: "list", lpar: vec![], @@ -888,7 +824,7 @@ pub(crate) fn fix_unnecessary_comprehension( })); } Expression::SetComp(inner) => { - body.value = Expression::Call(Box::new(Call { + tree = Expression::Call(Box::new(Call { func: Box::new(Expression::Name(Box::new(Name { value: "set", lpar: vec![], @@ -910,7 +846,7 @@ pub(crate) fn fix_unnecessary_comprehension( })); } Expression::DictComp(inner) => { - body.value = Expression::Call(Box::new(Call { + tree = Expression::Call(Box::new(Call { func: Box::new(Expression::Name(Box::new(Name { value: "dict", lpar: vec![], @@ -955,9 +891,8 @@ pub(crate) fn fix_unnecessary_map( kind: &str, ) -> Result { let module_text = locator.slice(expr.range()); - let mut tree = match_module(module_text)?; - let mut body = match_expr(&mut tree)?; - let call = match_call(body)?; + let mut tree = match_expression(module_text)?; + let call = match_call_mut(&mut tree)?; let arg = match_arg(call)?; let (args, lambda_func) = match &arg.value { @@ -976,9 +911,7 @@ pub(crate) fn fix_unnecessary_map( } }; - let Expression::Lambda(func_body) = &lambda_func else { - bail!("Expected a lambda") - }; + let func_body = match_lambda(&lambda_func)?; if args.len() == 2 { if func_body.params.params.iter().any(|f| f.default.is_some()) { @@ -1017,7 +950,7 @@ pub(crate) fn fix_unnecessary_map( match kind { "generator" => { - body.value = Expression::GeneratorExp(Box::new(GeneratorExp { + tree = Expression::GeneratorExp(Box::new(GeneratorExp { elt: func_body.body.clone(), for_in: compfor, lpar: vec![LeftParen::default()], @@ -1025,7 +958,7 @@ pub(crate) fn fix_unnecessary_map( })); } "list" => { - body.value = Expression::ListComp(Box::new(ListComp { + tree = Expression::ListComp(Box::new(ListComp { elt: func_body.body.clone(), for_in: compfor, lbracket: LeftSquareBracket::default(), @@ -1035,7 +968,7 @@ pub(crate) fn fix_unnecessary_map( })); } "set" => { - body.value = Expression::SetComp(Box::new(SetComp { + tree = Expression::SetComp(Box::new(SetComp { elt: func_body.body.clone(), for_in: compfor, lpar: vec![], @@ -1066,7 +999,7 @@ pub(crate) fn fix_unnecessary_map( bail!("Expected tuple for dict comprehension") }; - body.value = Expression::DictComp(Box::new(DictComp { + tree = Expression::DictComp(Box::new(DictComp { for_in: compfor, lpar: vec![], rpar: vec![], @@ -1115,12 +1048,11 @@ pub(crate) fn fix_unnecessary_literal_within_dict_call( expr: &rustpython_parser::ast::Expr, ) -> Result { let module_text = locator.slice(expr.range()); - let mut tree = match_module(module_text)?; - let mut body = match_expr(&mut tree)?; - let call = match_call(body)?; + let mut tree = match_expression(module_text)?; + let call = match_call_mut(&mut tree)?; let arg = match_arg(call)?; - body.value = arg.value.clone(); + tree = arg.value.clone(); let mut state = CodegenState { default_newline: &stylist.line_ending(), @@ -1140,9 +1072,8 @@ pub(crate) fn fix_unnecessary_comprehension_any_all( ) -> Result { // Expr(ListComp) -> Expr(GeneratorExp) let module_text = locator.slice(expr.range()); - let mut tree = match_module(module_text)?; - let body = match_expr(&mut tree)?; - let call = match_call(body)?; + let mut tree = match_expression(module_text)?; + let call = match_call_mut(&mut tree)?; let Expression::ListComp(list_comp) = &call.args[0].value else { bail!( diff --git a/crates/ruff/src/rules/flake8_pytest_style/rules/assertion.rs b/crates/ruff/src/rules/flake8_pytest_style/rules/assertion.rs index fecf968fcf..37fe645c19 100644 --- a/crates/ruff/src/rules/flake8_pytest_style/rules/assertion.rs +++ b/crates/ruff/src/rules/flake8_pytest_style/rules/assertion.rs @@ -5,7 +5,7 @@ use anyhow::Result; use libcst_native::{ Assert, BooleanOp, Codegen, CodegenState, CompoundStatement, Expression, ParenthesizableWhitespace, ParenthesizedNode, SimpleStatementLine, SimpleWhitespace, - SmallStatement, Statement, Suite, TrailingWhitespace, UnaryOp, UnaryOperation, + SmallStatement, Statement, TrailingWhitespace, UnaryOp, UnaryOperation, }; use rustpython_parser::ast::{self, Boolop, Excepthandler, Expr, Keyword, Ranged, Stmt, Unaryop}; @@ -17,6 +17,7 @@ use ruff_python_ast::visitor::Visitor; use ruff_python_ast::{visitor, whitespace}; use crate::checkers::ast::Checker; +use crate::cst::matchers::match_indented_block; use crate::cst::matchers::match_module; use crate::registry::AsRule; @@ -345,9 +346,7 @@ fn fix_composite_condition(stmt: &Stmt, locator: &Locator, stylist: &Stylist) -> bail!("Expected statement to be embedded in a function definition") }; - let Suite::IndentedBlock(indented_block) = &mut embedding.body else { - bail!("Expected indented block") - }; + let indented_block = match_indented_block(&mut embedding.body)?; indented_block.indent = Some(outer_indent); &mut indented_block.body diff --git a/crates/ruff/src/rules/flake8_simplify/rules/fix_if.rs b/crates/ruff/src/rules/flake8_simplify/rules/fix_if.rs index 7081dc69be..4a58149ff8 100644 --- a/crates/ruff/src/rules/flake8_simplify/rules/fix_if.rs +++ b/crates/ruff/src/rules/flake8_simplify/rules/fix_if.rs @@ -12,7 +12,7 @@ use ruff_diagnostics::Edit; use ruff_python_ast::source_code::{Locator, Stylist}; use ruff_python_ast::whitespace; -use crate::cst::matchers::match_module; +use crate::cst::matchers::{match_function_def, match_if, match_indented_block, match_statement}; fn parenthesize_and_operand(expr: Expression) -> Expression { match &expr { @@ -67,26 +67,23 @@ pub(crate) fn fix_nested_if_statements( }; // Parse the CST. - let mut tree = match_module(&module_text)?; + let mut tree = match_statement(&module_text)?; - let statements = if outer_indent.is_empty() { - &mut *tree.body + let statement = if outer_indent.is_empty() { + &mut tree } else { - let [Statement::Compound(CompoundStatement::FunctionDef(embedding))] = &mut *tree.body else { - bail!("Expected statement to be embedded in a function definition") - }; + let embedding = match_function_def(&mut tree)?; - let Suite::IndentedBlock(indented_block) = &mut embedding.body else { - bail!("Expected indented block") - }; + let indented_block = match_indented_block(&mut embedding.body)?; indented_block.indent = Some(outer_indent); - &mut *indented_block.body + let Some(statement) = indented_block.body.first_mut() else { + bail!("Expected indented block to have at least one statement") + }; + statement }; - let [Statement::Compound(CompoundStatement::If(outer_if))] = statements else { - bail!("Expected one outer if statement") - }; + let outer_if = match_if(statement)?; let If { body: Suite::IndentedBlock(ref mut outer_body), diff --git a/crates/ruff/src/rules/flake8_simplify/rules/fix_with.rs b/crates/ruff/src/rules/flake8_simplify/rules/fix_with.rs index ba6bc1ebc4..eaae0f8575 100644 --- a/crates/ruff/src/rules/flake8_simplify/rules/fix_with.rs +++ b/crates/ruff/src/rules/flake8_simplify/rules/fix_with.rs @@ -6,7 +6,7 @@ use ruff_diagnostics::Edit; use ruff_python_ast::source_code::{Locator, Stylist}; use ruff_python_ast::whitespace; -use crate::cst::matchers::match_module; +use crate::cst::matchers::{match_function_def, match_indented_block, match_statement, match_with}; /// (SIM117) Convert `with a: with b:` to `with a, b:`. pub(crate) fn fix_multiple_with_statements( @@ -32,26 +32,23 @@ pub(crate) fn fix_multiple_with_statements( }; // Parse the CST. - let mut tree = match_module(&module_text)?; + let mut tree = match_statement(&module_text)?; - let statements = if outer_indent.is_empty() { - &mut *tree.body + let statement = if outer_indent.is_empty() { + &mut tree } else { - let [Statement::Compound(CompoundStatement::FunctionDef(embedding))] = &mut *tree.body else { - bail!("Expected statement to be embedded in a function definition") - }; + let embedding = match_function_def(&mut tree)?; - let Suite::IndentedBlock(indented_block) = &mut embedding.body else { - bail!("Expected indented block") - }; + let indented_block = match_indented_block(&mut embedding.body)?; indented_block.indent = Some(outer_indent); - &mut *indented_block.body + let Some(statement) = indented_block.body.first_mut() else { + bail!("Expected indented block to have at least one statement") + }; + statement }; - let [Statement::Compound(CompoundStatement::With(outer_with))] = statements else { - bail!("Expected one outer with statement") - }; + let outer_with = match_with(statement)?; let With { body: Suite::IndentedBlock(ref mut outer_body), diff --git a/crates/ruff/src/rules/flake8_simplify/rules/key_in_dict.rs b/crates/ruff/src/rules/flake8_simplify/rules/key_in_dict.rs index d986ecd9b4..e63827a352 100644 --- a/crates/ruff/src/rules/flake8_simplify/rules/key_in_dict.rs +++ b/crates/ruff/src/rules/flake8_simplify/rules/key_in_dict.rs @@ -10,7 +10,7 @@ use ruff_macros::{derive_message_formats, violation}; use ruff_python_ast::source_code::{Locator, Stylist}; use crate::checkers::ast::Checker; -use crate::cst::matchers::{match_attribute, match_call, match_expression}; +use crate::cst::matchers::{match_attribute, match_call_mut, match_expression}; use crate::registry::AsRule; #[violation] @@ -39,7 +39,7 @@ fn get_value_content_for_key_in_dict( ) -> Result { let content = locator.slice(expr.range()); let mut expression = match_expression(content)?; - let call = match_call(&mut expression)?; + let call = match_call_mut(&mut expression)?; let attribute = match_attribute(&mut call.func)?; let mut state = CodegenState { diff --git a/crates/ruff/src/rules/pyflakes/fixes.rs b/crates/ruff/src/rules/pyflakes/fixes.rs index bc9ce26b33..6bc0e2aae3 100644 --- a/crates/ruff/src/rules/pyflakes/fixes.rs +++ b/crates/ruff/src/rules/pyflakes/fixes.rs @@ -12,7 +12,7 @@ use ruff_python_ast::source_code::{Locator, Stylist}; use ruff_python_ast::str::raw_contents; use crate::cst::matchers::{ - match_attribute, match_call, match_dict, match_expression, match_simple_string, + match_attribute, match_call_mut, match_dict, match_expression, match_simple_string, }; /// Generate a [`Edit`] to remove unused keys from format dict. @@ -52,7 +52,7 @@ pub(crate) fn remove_unused_keyword_arguments_from_format_call( ) -> Result { let module_text = locator.slice(location); let mut tree = match_expression(module_text)?; - let call = match_call(&mut tree)?; + let call = match_call_mut(&mut tree)?; call.args .retain(|e| !matches!(&e.keyword, Some(kw) if unused_arguments.contains(&kw.value))); @@ -135,7 +135,7 @@ pub(crate) fn remove_unused_positional_arguments_from_format_call( ) -> Result { let module_text = locator.slice(location); let mut tree = match_expression(module_text)?; - let call = match_call(&mut tree)?; + let call = match_call_mut(&mut tree)?; let mut index = 0; call.args.retain(|_| { diff --git a/crates/ruff/src/rules/pyupgrade/fixes.rs b/crates/ruff/src/rules/pyupgrade/fixes.rs index d29555109a..ce242175f1 100644 --- a/crates/ruff/src/rules/pyupgrade/fixes.rs +++ b/crates/ruff/src/rules/pyupgrade/fixes.rs @@ -1,8 +1,5 @@ -use anyhow::{bail, Result}; -use libcst_native::{ - Codegen, CodegenState, CompoundStatement, Expression, ParenthesizableWhitespace, - SmallStatement, Statement, Suite, -}; +use anyhow::Result; +use libcst_native::{Codegen, CodegenState, ParenthesizableWhitespace}; use ruff_text_size::{TextRange, TextSize}; use rustpython_parser::ast::{Expr, Ranged}; use rustpython_parser::{lexer, Mode, Tok}; @@ -10,7 +7,9 @@ use rustpython_parser::{lexer, Mode, Tok}; use ruff_diagnostics::Edit; use ruff_python_ast::source_code::{Locator, Stylist}; -use crate::cst::matchers::match_module; +use crate::cst::matchers::{ + match_call_mut, match_expression, match_function_def, match_indented_block, match_statement, +}; /// Safely adjust the indentation of the indented block at [`TextRange`]. pub(crate) fn adjust_indentation( @@ -23,15 +22,11 @@ pub(crate) fn adjust_indentation( let module_text = format!("def f():{}{contents}", stylist.line_ending().as_str()); - let mut tree = match_module(&module_text)?; + let mut tree = match_statement(&module_text)?; - let [Statement::Compound(CompoundStatement::FunctionDef(embedding))] = &mut *tree.body else { - bail!("Expected statement to be embedded in a function definition") - }; + let embedding = match_function_def(&mut tree)?; - let Suite::IndentedBlock(indented_block) = &mut embedding.body else { - bail!("Expected indented block") - }; + let indented_block = match_indented_block(&mut embedding.body)?; indented_block.indent = Some(indentation); let mut state = CodegenState { @@ -58,17 +53,9 @@ pub(crate) fn remove_super_arguments( let range = expr.range(); let contents = locator.slice(range); - let mut tree = libcst_native::parse_module(contents, None).ok()?; + let mut tree = match_expression(contents).ok()?; - let Statement::Simple(body) = tree.body.first_mut()? else { - return None; - }; - let SmallStatement::Expr(body) = body.body.first_mut()? else { - return None; - }; - let Expression::Call(body) = &mut body.value else { - return None; - }; + let body = match_call_mut(&mut tree).ok()?; body.args = vec![]; body.whitespace_before_args = ParenthesizableWhitespace::default(); diff --git a/crates/ruff/src/rules/pyupgrade/rules/deprecated_mock_import.rs b/crates/ruff/src/rules/pyupgrade/rules/deprecated_mock_import.rs index 6745134fd4..d99c272926 100644 --- a/crates/ruff/src/rules/pyupgrade/rules/deprecated_mock_import.rs +++ b/crates/ruff/src/rules/pyupgrade/rules/deprecated_mock_import.rs @@ -13,7 +13,7 @@ use ruff_python_ast::source_code::{Locator, Stylist}; use ruff_python_ast::whitespace::indentation; use crate::checkers::ast::Checker; -use crate::cst::matchers::{match_import, match_import_from, match_module}; +use crate::cst::matchers::{match_import, match_import_from, match_statement}; use crate::registry::{AsRule, Rule}; #[derive(Debug, PartialEq, Eq, Copy, Clone)] @@ -126,7 +126,7 @@ fn format_import( stylist: &Stylist, ) -> Result { let module_text = locator.slice(stmt.range()); - let mut tree = match_module(module_text)?; + let mut tree = match_statement(module_text)?; let mut import = match_import(&mut tree)?; let Import { names, .. } = import.clone(); @@ -160,7 +160,7 @@ fn format_import_from( stylist: &Stylist, ) -> Result { let module_text = locator.slice(stmt.range()); - let mut tree = match_module(module_text).unwrap(); + let mut tree = match_statement(module_text).unwrap(); let mut import = match_import_from(&mut tree)?; if let ImportFrom { diff --git a/crates/ruff/src/rules/pyupgrade/rules/format_literals.rs b/crates/ruff/src/rules/pyupgrade/rules/format_literals.rs index 802d63d91f..5352f3c008 100644 --- a/crates/ruff/src/rules/pyupgrade/rules/format_literals.rs +++ b/crates/ruff/src/rules/pyupgrade/rules/format_literals.rs @@ -1,5 +1,5 @@ use anyhow::{anyhow, bail, Result}; -use libcst_native::{Arg, Codegen, CodegenState, Expression}; +use libcst_native::{Arg, Codegen, CodegenState}; use once_cell::sync::Lazy; use regex::Regex; use rustpython_parser::ast::{Expr, Ranged}; @@ -9,7 +9,7 @@ use ruff_macros::{derive_message_formats, violation}; use ruff_python_ast::source_code::{Locator, Stylist}; use crate::checkers::ast::Checker; -use crate::cst::matchers::{match_call, match_expression}; +use crate::cst::matchers::{match_attribute, match_call_mut, match_expression}; use crate::registry::AsRule; use crate::rules::pyflakes::format::FormatSummary; @@ -89,7 +89,7 @@ fn generate_call( ) -> Result { let module_text = locator.slice(expr.range()); let mut expression = match_expression(module_text)?; - let mut call = match_call(&mut expression)?; + let mut call = match_call_mut(&mut expression)?; // Fix the call arguments. if !is_sequential(correct_order) { @@ -97,9 +97,7 @@ fn generate_call( } // Fix the string itself. - let Expression::Attribute(item) = &*call.func else { - panic!("Expected: Expression::Attribute") - }; + let item = match_attribute(&mut call.func)?; let mut state = CodegenState { default_newline: &stylist.line_ending(), diff --git a/crates/ruff/src/rules/ruff/rules/explicit_f_string_type_conversion.rs b/crates/ruff/src/rules/ruff/rules/explicit_f_string_type_conversion.rs index e281e791b1..a08071d29d 100644 --- a/crates/ruff/src/rules/ruff/rules/explicit_f_string_type_conversion.rs +++ b/crates/ruff/src/rules/ruff/rules/explicit_f_string_type_conversion.rs @@ -8,7 +8,7 @@ use ruff_python_ast::source_code::{Locator, Stylist}; use crate::checkers::ast::Checker; use crate::cst::matchers::{ - match_call, match_expression, match_formatted_string, match_formatted_string_expression, + match_call_mut, match_expression, match_formatted_string, match_formatted_string_expression, match_name, }; use crate::registry::AsRule; @@ -61,8 +61,8 @@ fn fix_explicit_f_string_type_conversion( // Replace the formatted call expression at `index` with a conversion flag. let mut formatted_string_expression = match_formatted_string_expression(&mut formatted_string.parts[index])?; - let call = match_call(&mut formatted_string_expression.expression)?; - let name = match_name(&mut call.func)?; + let call = match_call_mut(&mut formatted_string_expression.expression)?; + let name = match_name(&call.func)?; match name.value { "str" => { formatted_string_expression.conversion = Some("s");