Rework CST matchers (#4536)

Co-authored-by: Micha Reiser <micha@reiser.io>
This commit is contained in:
Jonathan Plasse
2023-05-23 08:26:51 +02:00
committed by GitHub
parent 74effb40b9
commit 1ddc577204
13 changed files with 250 additions and 245 deletions

View File

@@ -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<Edit> {
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");
};

View File

@@ -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<Module> {
@@ -19,20 +20,15 @@ pub(crate) fn match_expression(expression_text: &str) -> Result<Expression> {
}
}
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<Statement> {
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")
}
}

View File

@@ -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<Edit> {
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()))
}
}

View File

@@ -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<Edit> {
// 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<Edit> {
// 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<Edit> {
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<Edit> {
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<Edit> {
// 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<Edit> {
// 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::<Result<Vec<DictElement>>>()?;
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<Edit> {
// 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<Edit> {
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<Edit> {
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<Edit> {
// 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<Edit> {
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<Edit> {
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<Edit> {
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<Edit> {
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<Edit> {
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<Fix> {
// 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!(

View File

@@ -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

View File

@@ -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),

View File

@@ -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),

View File

@@ -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<String> {
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 {

View File

@@ -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<Edit> {
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<Edit> {
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(|_| {

View File

@@ -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();

View File

@@ -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<String> {
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<String> {
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 {

View File

@@ -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<String> {
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(),

View File

@@ -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");