mirror of https://github.com/astral-sh/ruff
1294 lines
42 KiB
Rust
1294 lines
42 KiB
Rust
use anyhow::{bail, Result};
|
|
use itertools::Itertools;
|
|
use libcst_native::{
|
|
Arg, AssignEqual, AssignTargetExpression, Call, Comment, CompFor, Dict, DictComp, DictElement,
|
|
Element, EmptyLine, Expression, GeneratorExp, LeftCurlyBrace, LeftParen, LeftSquareBracket,
|
|
List, ListComp, Name, ParenthesizableWhitespace, ParenthesizedNode, ParenthesizedWhitespace,
|
|
RightCurlyBrace, RightParen, RightSquareBracket, Set, SetComp, SimpleString, SimpleWhitespace,
|
|
TrailingWhitespace, Tuple,
|
|
};
|
|
|
|
use ruff_diagnostics::{Edit, Fix};
|
|
use ruff_python_ast::Expr;
|
|
use ruff_python_codegen::Stylist;
|
|
use ruff_python_semantic::SemanticModel;
|
|
use ruff_source_file::Locator;
|
|
use ruff_text_size::{Ranged, TextRange};
|
|
|
|
use crate::cst::helpers::{negate, space};
|
|
use crate::fix::codemods::CodegenStylist;
|
|
use crate::fix::edits::pad;
|
|
use crate::rules::flake8_comprehensions::rules::ObjectType;
|
|
use crate::{
|
|
checkers::ast::Checker,
|
|
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(
|
|
expr: &Expr,
|
|
locator: &Locator,
|
|
stylist: &Stylist,
|
|
) -> Result<Edit> {
|
|
// Expr(Call(GeneratorExp)))) -> Expr(ListComp)))
|
|
let module_text = locator.slice(expr);
|
|
let mut tree = match_expression(module_text)?;
|
|
let call = match_call_mut(&mut tree)?;
|
|
let arg = match_arg(call)?;
|
|
|
|
let generator_exp = match_generator_exp(&arg.value)?;
|
|
|
|
tree = Expression::ListComp(Box::new(ListComp {
|
|
elt: generator_exp.elt.clone(),
|
|
for_in: generator_exp.for_in.clone(),
|
|
lbracket: LeftSquareBracket {
|
|
whitespace_after: call.whitespace_before_args.clone(),
|
|
},
|
|
rbracket: RightSquareBracket {
|
|
whitespace_before: arg.whitespace_after_arg.clone(),
|
|
},
|
|
lpar: generator_exp.lpar.clone(),
|
|
rpar: generator_exp.rpar.clone(),
|
|
}));
|
|
|
|
Ok(Edit::range_replacement(
|
|
tree.codegen_stylist(stylist),
|
|
expr.range(),
|
|
))
|
|
}
|
|
|
|
/// (C401) Convert `set(x for x in y)` to `{x for x in y}`.
|
|
pub(crate) fn fix_unnecessary_generator_set(expr: &Expr, checker: &Checker) -> Result<Edit> {
|
|
let locator = checker.locator();
|
|
let stylist = checker.stylist();
|
|
|
|
// Expr(Call(GeneratorExp)))) -> Expr(SetComp)))
|
|
let module_text = locator.slice(expr);
|
|
let mut tree = match_expression(module_text)?;
|
|
let call = match_call_mut(&mut tree)?;
|
|
let arg = match_arg(call)?;
|
|
|
|
let generator_exp = match_generator_exp(&arg.value)?;
|
|
|
|
tree = Expression::SetComp(Box::new(SetComp {
|
|
elt: generator_exp.elt.clone(),
|
|
for_in: generator_exp.for_in.clone(),
|
|
lbrace: LeftCurlyBrace {
|
|
whitespace_after: call.whitespace_before_args.clone(),
|
|
},
|
|
rbrace: RightCurlyBrace {
|
|
whitespace_before: arg.whitespace_after_arg.clone(),
|
|
},
|
|
lpar: generator_exp.lpar.clone(),
|
|
rpar: generator_exp.rpar.clone(),
|
|
}));
|
|
|
|
let content = tree.codegen_stylist(stylist);
|
|
|
|
Ok(Edit::range_replacement(
|
|
pad_expression(content, expr.range(), checker.locator(), checker.semantic()),
|
|
expr.range(),
|
|
))
|
|
}
|
|
|
|
/// (C402) Convert `dict((x, x) for x in range(3))` to `{x: x for x in
|
|
/// range(3)}`.
|
|
pub(crate) fn fix_unnecessary_generator_dict(expr: &Expr, checker: &Checker) -> Result<Edit> {
|
|
let locator = checker.locator();
|
|
let stylist = checker.stylist();
|
|
|
|
let module_text = locator.slice(expr);
|
|
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 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");
|
|
};
|
|
|
|
// Insert whitespace before the `for`, since we're removing parentheses, as in:
|
|
// ```python
|
|
// dict((x, x)for x in range(3))
|
|
// ```
|
|
let mut for_in = generator_exp.for_in.clone();
|
|
if for_in.whitespace_before == ParenthesizableWhitespace::default() {
|
|
for_in.whitespace_before =
|
|
ParenthesizableWhitespace::SimpleWhitespace(SimpleWhitespace(" "));
|
|
}
|
|
|
|
tree = Expression::DictComp(Box::new(DictComp {
|
|
key: Box::new(key.clone()),
|
|
value: Box::new(value.clone()),
|
|
for_in,
|
|
lbrace: LeftCurlyBrace {
|
|
whitespace_after: call.whitespace_before_args.clone(),
|
|
},
|
|
rbrace: RightCurlyBrace {
|
|
whitespace_before: arg.whitespace_after_arg.clone(),
|
|
},
|
|
lpar: vec![],
|
|
rpar: vec![],
|
|
whitespace_before_colon: ParenthesizableWhitespace::default(),
|
|
whitespace_after_colon: space(),
|
|
}));
|
|
|
|
Ok(Edit::range_replacement(
|
|
pad_expression(
|
|
tree.codegen_stylist(stylist),
|
|
expr.range(),
|
|
checker.locator(),
|
|
checker.semantic(),
|
|
),
|
|
expr.range(),
|
|
))
|
|
}
|
|
|
|
/// (C403) Convert `set([x for x in y])` to `{x for x in y}`.
|
|
pub(crate) fn fix_unnecessary_list_comprehension_set(
|
|
expr: &Expr,
|
|
checker: &Checker,
|
|
) -> Result<Edit> {
|
|
let locator = checker.locator();
|
|
let stylist = checker.stylist();
|
|
// Expr(Call(ListComp)))) ->
|
|
// Expr(SetComp)))
|
|
let module_text = locator.slice(expr);
|
|
let mut tree = match_expression(module_text)?;
|
|
let call = match_call_mut(&mut tree)?;
|
|
let arg = match_arg(call)?;
|
|
|
|
let list_comp = match_list_comp(&arg.value)?;
|
|
|
|
tree = Expression::SetComp(Box::new(SetComp {
|
|
elt: list_comp.elt.clone(),
|
|
for_in: list_comp.for_in.clone(),
|
|
lbrace: LeftCurlyBrace {
|
|
whitespace_after: call.whitespace_before_args.clone(),
|
|
},
|
|
rbrace: RightCurlyBrace {
|
|
whitespace_before: arg.whitespace_after_arg.clone(),
|
|
},
|
|
lpar: list_comp.lpar.clone(),
|
|
rpar: list_comp.rpar.clone(),
|
|
}));
|
|
|
|
Ok(Edit::range_replacement(
|
|
pad_expression(
|
|
tree.codegen_stylist(stylist),
|
|
expr.range(),
|
|
checker.locator(),
|
|
checker.semantic(),
|
|
),
|
|
expr.range(),
|
|
))
|
|
}
|
|
|
|
/// (C404) Convert `dict([(i, i) for i in range(3)])` to `{i: i for i in
|
|
/// range(3)}`.
|
|
pub(crate) fn fix_unnecessary_list_comprehension_dict(
|
|
expr: &Expr,
|
|
checker: &Checker,
|
|
) -> Result<Edit> {
|
|
let locator = checker.locator();
|
|
let stylist = checker.stylist();
|
|
|
|
let module_text = locator.slice(expr);
|
|
let mut tree = match_expression(module_text)?;
|
|
let call = match_call_mut(&mut tree)?;
|
|
let arg = match_arg(call)?;
|
|
|
|
let list_comp = match_list_comp(&arg.value)?;
|
|
|
|
let tuple = match_tuple(&list_comp.elt)?;
|
|
|
|
let [Element::Simple { value: key, .. }, Element::Simple { value, .. }] = &tuple.elements[..]
|
|
else {
|
|
bail!("Expected tuple with two elements");
|
|
};
|
|
|
|
// Insert whitespace before the `for`, since we're removing parentheses, as in:
|
|
// ```python
|
|
// dict((x, x)for x in range(3))
|
|
// ```
|
|
let mut for_in = list_comp.for_in.clone();
|
|
if for_in.whitespace_before == ParenthesizableWhitespace::default() {
|
|
for_in.whitespace_before =
|
|
ParenthesizableWhitespace::SimpleWhitespace(SimpleWhitespace(" "));
|
|
}
|
|
|
|
tree = Expression::DictComp(Box::new(DictComp {
|
|
key: Box::new(key.clone()),
|
|
value: Box::new(value.clone()),
|
|
for_in,
|
|
whitespace_before_colon: ParenthesizableWhitespace::default(),
|
|
whitespace_after_colon: space(),
|
|
lbrace: LeftCurlyBrace {
|
|
whitespace_after: call.whitespace_before_args.clone(),
|
|
},
|
|
rbrace: RightCurlyBrace {
|
|
whitespace_before: arg.whitespace_after_arg.clone(),
|
|
},
|
|
lpar: list_comp.lpar.clone(),
|
|
rpar: list_comp.rpar.clone(),
|
|
}));
|
|
|
|
Ok(Edit::range_replacement(
|
|
pad_expression(
|
|
tree.codegen_stylist(stylist),
|
|
expr.range(),
|
|
checker.locator(),
|
|
checker.semantic(),
|
|
),
|
|
expr.range(),
|
|
))
|
|
}
|
|
|
|
/// Drop a trailing comma from a list of tuple elements.
|
|
fn drop_trailing_comma<'a>(
|
|
tuple: &Tuple<'a>,
|
|
) -> Result<(
|
|
Vec<Element<'a>>,
|
|
ParenthesizableWhitespace<'a>,
|
|
ParenthesizableWhitespace<'a>,
|
|
)> {
|
|
let whitespace_after = tuple
|
|
.lpar
|
|
.first()
|
|
.ok_or_else(|| anyhow::anyhow!("Expected at least one set of parentheses"))?
|
|
.whitespace_after
|
|
.clone();
|
|
let whitespace_before = tuple
|
|
.rpar
|
|
.first()
|
|
.ok_or_else(|| anyhow::anyhow!("Expected at least one set of parentheses"))?
|
|
.whitespace_before
|
|
.clone();
|
|
|
|
let mut elements = tuple.elements.clone();
|
|
if elements.len() == 1 {
|
|
if let Some(Element::Simple {
|
|
value,
|
|
comma: Some(..),
|
|
..
|
|
}) = elements.last()
|
|
{
|
|
if whitespace_before == ParenthesizableWhitespace::default()
|
|
&& whitespace_after == ParenthesizableWhitespace::default()
|
|
{
|
|
elements[0] = Element::Simple {
|
|
value: value.clone(),
|
|
comma: None,
|
|
};
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok((elements, whitespace_after, whitespace_before))
|
|
}
|
|
|
|
/// (C405) Convert `set((1, 2))` to `{1, 2}`.
|
|
pub(crate) fn fix_unnecessary_literal_set(expr: &Expr, checker: &Checker) -> Result<Edit> {
|
|
let locator = checker.locator();
|
|
let stylist = checker.stylist();
|
|
|
|
// Expr(Call(List|Tuple)))) -> Expr(Set)))
|
|
let module_text = locator.slice(expr);
|
|
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) => drop_trailing_comma(inner)?,
|
|
Expression::List(inner) => (
|
|
inner.elements.clone(),
|
|
inner.lbracket.whitespace_after.clone(),
|
|
inner.rbracket.whitespace_before.clone(),
|
|
),
|
|
_ => {
|
|
bail!("Expected Expression::Tuple | Expression::List");
|
|
}
|
|
};
|
|
|
|
if elements.is_empty() {
|
|
call.args = vec![];
|
|
} else {
|
|
tree = Expression::Set(Box::new(Set {
|
|
elements,
|
|
lbrace: LeftCurlyBrace { whitespace_after },
|
|
rbrace: RightCurlyBrace { whitespace_before },
|
|
lpar: vec![],
|
|
rpar: vec![],
|
|
}));
|
|
}
|
|
|
|
Ok(Edit::range_replacement(
|
|
pad_expression(
|
|
tree.codegen_stylist(stylist),
|
|
expr.range(),
|
|
checker.locator(),
|
|
checker.semantic(),
|
|
),
|
|
expr.range(),
|
|
))
|
|
}
|
|
|
|
/// (C406) Convert `dict([(1, 2)])` to `{1: 2}`.
|
|
pub(crate) fn fix_unnecessary_literal_dict(expr: &Expr, checker: &Checker) -> Result<Edit> {
|
|
let locator = checker.locator();
|
|
let stylist = checker.stylist();
|
|
|
|
// Expr(Call(List|Tuple)))) -> Expr(Dict)))
|
|
let module_text = locator.slice(expr);
|
|
let mut tree = match_expression(module_text)?;
|
|
let call = match_call_mut(&mut tree)?;
|
|
let arg = match_arg(call)?;
|
|
|
|
let elements = match &arg.value {
|
|
Expression::Tuple(inner) => &inner.elements,
|
|
Expression::List(inner) => &inner.elements,
|
|
_ => {
|
|
bail!("Expected Expression::Tuple | Expression::List");
|
|
}
|
|
};
|
|
|
|
let elements: Vec<DictElement> = elements
|
|
.iter()
|
|
.map(|element| {
|
|
if let Element::Simple {
|
|
value: Expression::Tuple(tuple),
|
|
comma,
|
|
} = element
|
|
{
|
|
if let Some(Element::Simple { value: key, .. }) = tuple.elements.get(0) {
|
|
if let Some(Element::Simple { value, .. }) = tuple.elements.get(1) {
|
|
return Ok(DictElement::Simple {
|
|
key: key.clone(),
|
|
value: value.clone(),
|
|
comma: comma.clone(),
|
|
whitespace_before_colon: ParenthesizableWhitespace::default(),
|
|
whitespace_after_colon: ParenthesizableWhitespace::SimpleWhitespace(
|
|
SimpleWhitespace(" "),
|
|
),
|
|
});
|
|
}
|
|
}
|
|
}
|
|
bail!("Expected each argument to be a tuple of length two")
|
|
})
|
|
.collect::<Result<Vec<DictElement>>>()?;
|
|
|
|
tree = Expression::Dict(Box::new(Dict {
|
|
elements,
|
|
lbrace: LeftCurlyBrace {
|
|
whitespace_after: call.whitespace_before_args.clone(),
|
|
},
|
|
rbrace: RightCurlyBrace {
|
|
whitespace_before: arg.whitespace_after_arg.clone(),
|
|
},
|
|
lpar: vec![],
|
|
rpar: vec![],
|
|
}));
|
|
|
|
Ok(Edit::range_replacement(
|
|
pad_expression(
|
|
tree.codegen_stylist(stylist),
|
|
expr.range(),
|
|
checker.locator(),
|
|
checker.semantic(),
|
|
),
|
|
expr.range(),
|
|
))
|
|
}
|
|
|
|
/// (C408)
|
|
pub(crate) fn fix_unnecessary_collection_call(expr: &Expr, checker: &Checker) -> Result<Edit> {
|
|
enum Collection {
|
|
Tuple,
|
|
List,
|
|
Dict,
|
|
}
|
|
|
|
let locator = checker.locator();
|
|
let stylist = checker.stylist();
|
|
|
|
// Expr(Call("list" | "tuple" | "dict")))) -> Expr(List|Tuple|Dict)
|
|
let module_text = locator.slice(expr);
|
|
let mut tree = match_expression(module_text)?;
|
|
let call = match_call(&tree)?;
|
|
let name = match_name(&call.func)?;
|
|
let collection = match name.value {
|
|
"tuple" => Collection::Tuple,
|
|
"list" => Collection::List,
|
|
"dict" => Collection::Dict,
|
|
_ => bail!("Expected 'tuple', 'list', or 'dict'"),
|
|
};
|
|
|
|
// Arena allocator used to create formatted strings of sufficient lifetime,
|
|
// below.
|
|
let mut arena: Vec<String> = vec![];
|
|
|
|
match collection {
|
|
Collection::Tuple => {
|
|
tree = Expression::Tuple(Box::new(Tuple {
|
|
elements: vec![],
|
|
lpar: vec![LeftParen::default()],
|
|
rpar: vec![RightParen::default()],
|
|
}));
|
|
}
|
|
Collection::List => {
|
|
tree = Expression::List(Box::new(List {
|
|
elements: vec![],
|
|
lbracket: LeftSquareBracket::default(),
|
|
rbracket: RightSquareBracket::default(),
|
|
lpar: vec![],
|
|
rpar: vec![],
|
|
}));
|
|
}
|
|
Collection::Dict => {
|
|
if call.args.is_empty() {
|
|
tree = Expression::Dict(Box::new(Dict {
|
|
elements: vec![],
|
|
lbrace: LeftCurlyBrace::default(),
|
|
rbrace: RightCurlyBrace::default(),
|
|
lpar: vec![],
|
|
rpar: vec![],
|
|
}));
|
|
} else {
|
|
let quote = checker.f_string_quote_style().unwrap_or(stylist.quote());
|
|
|
|
// Quote each argument.
|
|
for arg in &call.args {
|
|
let quoted = format!(
|
|
"{}{}{}",
|
|
quote,
|
|
arg.keyword
|
|
.as_ref()
|
|
.expect("Expected dictionary argument to be kwarg")
|
|
.value,
|
|
quote,
|
|
);
|
|
arena.push(quoted);
|
|
}
|
|
|
|
let elements = call
|
|
.args
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, arg)| DictElement::Simple {
|
|
key: Expression::SimpleString(Box::new(SimpleString {
|
|
value: &arena[i],
|
|
lpar: vec![],
|
|
rpar: vec![],
|
|
})),
|
|
value: arg.value.clone(),
|
|
comma: arg.comma.clone(),
|
|
whitespace_before_colon: ParenthesizableWhitespace::default(),
|
|
whitespace_after_colon: ParenthesizableWhitespace::SimpleWhitespace(
|
|
SimpleWhitespace(" "),
|
|
),
|
|
})
|
|
.collect();
|
|
|
|
tree = Expression::Dict(Box::new(Dict {
|
|
elements,
|
|
lbrace: LeftCurlyBrace {
|
|
whitespace_after: call.whitespace_before_args.clone(),
|
|
},
|
|
rbrace: RightCurlyBrace {
|
|
whitespace_before: call
|
|
.args
|
|
.last()
|
|
.expect("Arguments should be non-empty")
|
|
.whitespace_after_arg
|
|
.clone(),
|
|
},
|
|
lpar: vec![],
|
|
rpar: vec![],
|
|
}));
|
|
}
|
|
}
|
|
};
|
|
|
|
Ok(Edit::range_replacement(
|
|
if matches!(collection, Collection::Dict) {
|
|
pad_expression(
|
|
tree.codegen_stylist(stylist),
|
|
expr.range(),
|
|
checker.locator(),
|
|
checker.semantic(),
|
|
)
|
|
} else {
|
|
tree.codegen_stylist(stylist)
|
|
},
|
|
expr.range(),
|
|
))
|
|
}
|
|
|
|
/// Re-formats the given expression for use within a formatted string.
|
|
///
|
|
/// For example, when converting a `dict` call to a dictionary literal within
|
|
/// a formatted string, we might naively generate the following code:
|
|
///
|
|
/// ```python
|
|
/// f"{{'a': 1, 'b': 2}}"
|
|
/// ```
|
|
///
|
|
/// However, this is a syntax error under the f-string grammar. As such,
|
|
/// this method will pad the start and end of an expression as needed to
|
|
/// avoid producing invalid syntax.
|
|
fn pad_expression(
|
|
content: String,
|
|
range: TextRange,
|
|
locator: &Locator,
|
|
semantic: &SemanticModel,
|
|
) -> String {
|
|
if !semantic.in_f_string() {
|
|
return content;
|
|
}
|
|
|
|
// If the expression is immediately preceded by an opening brace, then
|
|
// we need to add a space before the expression.
|
|
let prefix = locator.up_to(range.start());
|
|
let left_pad = matches!(prefix.chars().next_back(), Some('{'));
|
|
|
|
// If the expression is immediately preceded by an opening brace, then
|
|
// we need to add a space before the expression.
|
|
let suffix = locator.after(range.end());
|
|
let right_pad = matches!(suffix.chars().next(), Some('}'));
|
|
|
|
if left_pad && right_pad {
|
|
format!(" {content} ")
|
|
} else if left_pad {
|
|
format!(" {content}")
|
|
} else if right_pad {
|
|
format!("{content} ")
|
|
} else {
|
|
content
|
|
}
|
|
}
|
|
|
|
/// (C409) Convert `tuple([1, 2])` to `tuple(1, 2)`
|
|
pub(crate) fn fix_unnecessary_literal_within_tuple_call(
|
|
expr: &Expr,
|
|
locator: &Locator,
|
|
stylist: &Stylist,
|
|
) -> Result<Edit> {
|
|
let module_text = locator.slice(expr);
|
|
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) => (
|
|
&inner.elements,
|
|
&inner
|
|
.lpar
|
|
.first()
|
|
.ok_or_else(|| anyhow::anyhow!("Expected at least one set of parentheses"))?
|
|
.whitespace_after,
|
|
&inner
|
|
.rpar
|
|
.first()
|
|
.ok_or_else(|| anyhow::anyhow!("Expected at least one set of parentheses"))?
|
|
.whitespace_before,
|
|
),
|
|
Expression::List(inner) => (
|
|
&inner.elements,
|
|
&inner.lbracket.whitespace_after,
|
|
&inner.rbracket.whitespace_before,
|
|
),
|
|
_ => {
|
|
bail!("Expected Expression::Tuple | Expression::List");
|
|
}
|
|
};
|
|
|
|
tree = Expression::Tuple(Box::new(Tuple {
|
|
elements: elements.clone(),
|
|
lpar: vec![LeftParen {
|
|
whitespace_after: whitespace_after.clone(),
|
|
}],
|
|
rpar: vec![RightParen {
|
|
whitespace_before: whitespace_before.clone(),
|
|
}],
|
|
}));
|
|
|
|
Ok(Edit::range_replacement(
|
|
tree.codegen_stylist(stylist),
|
|
expr.range(),
|
|
))
|
|
}
|
|
|
|
/// (C410) Convert `list([1, 2])` to `[1, 2]`
|
|
pub(crate) fn fix_unnecessary_literal_within_list_call(
|
|
expr: &Expr,
|
|
locator: &Locator,
|
|
stylist: &Stylist,
|
|
) -> Result<Edit> {
|
|
let module_text = locator.slice(expr);
|
|
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) => (
|
|
&inner.elements,
|
|
&inner
|
|
.lpar
|
|
.first()
|
|
.ok_or_else(|| anyhow::anyhow!("Expected at least one set of parentheses"))?
|
|
.whitespace_after,
|
|
&inner
|
|
.rpar
|
|
.first()
|
|
.ok_or_else(|| anyhow::anyhow!("Expected at least one set of parentheses"))?
|
|
.whitespace_before,
|
|
),
|
|
Expression::List(inner) => (
|
|
&inner.elements,
|
|
&inner.lbracket.whitespace_after,
|
|
&inner.rbracket.whitespace_before,
|
|
),
|
|
_ => {
|
|
bail!("Expected Expression::Tuple | Expression::List");
|
|
}
|
|
};
|
|
|
|
tree = Expression::List(Box::new(List {
|
|
elements: elements.clone(),
|
|
lbracket: LeftSquareBracket {
|
|
whitespace_after: whitespace_after.clone(),
|
|
},
|
|
rbracket: RightSquareBracket {
|
|
whitespace_before: whitespace_before.clone(),
|
|
},
|
|
lpar: vec![],
|
|
rpar: vec![],
|
|
}));
|
|
|
|
Ok(Edit::range_replacement(
|
|
tree.codegen_stylist(stylist),
|
|
expr.range(),
|
|
))
|
|
}
|
|
|
|
/// (C411) Convert `list([i * i for i in x])` to `[i * i for i in x]`.
|
|
pub(crate) fn fix_unnecessary_list_call(
|
|
expr: &Expr,
|
|
locator: &Locator,
|
|
stylist: &Stylist,
|
|
) -> Result<Edit> {
|
|
// Expr(Call(List|Tuple)))) -> Expr(List|Tuple)))
|
|
let module_text = locator.slice(expr);
|
|
let mut tree = match_expression(module_text)?;
|
|
let call = match_call_mut(&mut tree)?;
|
|
let arg = match_arg(call)?;
|
|
|
|
tree = arg.value.clone();
|
|
|
|
Ok(Edit::range_replacement(
|
|
tree.codegen_stylist(stylist),
|
|
expr.range(),
|
|
))
|
|
}
|
|
|
|
/// (C413) Convert `list(sorted([2, 3, 1]))` to `sorted([2, 3, 1])`.
|
|
/// (C413) Convert `reversed(sorted([2, 3, 1]))` to `sorted([2, 3, 1],
|
|
/// reverse=True)`.
|
|
pub(crate) fn fix_unnecessary_call_around_sorted(
|
|
expr: &Expr,
|
|
locator: &Locator,
|
|
stylist: &Stylist,
|
|
) -> Result<Edit> {
|
|
let module_text = locator.slice(expr);
|
|
let mut tree = match_expression(module_text)?;
|
|
let outer_call = match_call_mut(&mut tree)?;
|
|
let inner_call = match &outer_call.args[..] {
|
|
[arg] => match_call(&arg.value)?,
|
|
_ => {
|
|
bail!("Expected one argument in outer function call");
|
|
}
|
|
};
|
|
|
|
if let Expression::Name(outer_name) = &*outer_call.func {
|
|
if outer_name.value == "list" {
|
|
tree = Expression::Call(Box::new((*inner_call).clone()));
|
|
} else {
|
|
// If the `reverse` argument is used...
|
|
let args = if inner_call.args.iter().any(|arg| {
|
|
matches!(
|
|
arg.keyword,
|
|
Some(Name {
|
|
value: "reverse",
|
|
..
|
|
})
|
|
)
|
|
}) {
|
|
// Negate the `reverse` argument.
|
|
inner_call
|
|
.args
|
|
.clone()
|
|
.into_iter()
|
|
.map(|mut arg| {
|
|
if matches!(
|
|
arg.keyword,
|
|
Some(Name {
|
|
value: "reverse",
|
|
..
|
|
})
|
|
) {
|
|
arg.value = negate(&arg.value);
|
|
}
|
|
arg
|
|
})
|
|
.collect_vec()
|
|
} else {
|
|
let mut args = inner_call.args.clone();
|
|
|
|
// If necessary, parenthesize a generator expression, as a generator expression must
|
|
// be parenthesized if it's not a solitary argument. For example, given:
|
|
// ```python
|
|
// reversed(sorted(i for i in range(42)))
|
|
// ```
|
|
// Rewrite as:
|
|
// ```python
|
|
// sorted((i for i in range(42)), reverse=True)
|
|
// ```
|
|
if let [arg] = args.as_mut_slice() {
|
|
if matches!(arg.value, Expression::GeneratorExp(_)) {
|
|
if arg.value.lpar().is_empty() && arg.value.rpar().is_empty() {
|
|
arg.value = arg
|
|
.value
|
|
.clone()
|
|
.with_parens(LeftParen::default(), RightParen::default());
|
|
}
|
|
}
|
|
}
|
|
|
|
// Add the `reverse=True` argument.
|
|
args.push(Arg {
|
|
value: Expression::Name(Box::new(Name {
|
|
value: "True",
|
|
lpar: vec![],
|
|
rpar: vec![],
|
|
})),
|
|
keyword: Some(Name {
|
|
value: "reverse",
|
|
lpar: vec![],
|
|
rpar: vec![],
|
|
}),
|
|
equal: Some(AssignEqual {
|
|
whitespace_before: ParenthesizableWhitespace::default(),
|
|
whitespace_after: ParenthesizableWhitespace::default(),
|
|
}),
|
|
comma: None,
|
|
star: "",
|
|
whitespace_after_star: ParenthesizableWhitespace::default(),
|
|
whitespace_after_arg: ParenthesizableWhitespace::default(),
|
|
});
|
|
args
|
|
};
|
|
|
|
tree = Expression::Call(Box::new(Call {
|
|
func: inner_call.func.clone(),
|
|
args,
|
|
lpar: inner_call.lpar.clone(),
|
|
rpar: inner_call.rpar.clone(),
|
|
whitespace_after_func: inner_call.whitespace_after_func.clone(),
|
|
whitespace_before_args: inner_call.whitespace_before_args.clone(),
|
|
}));
|
|
}
|
|
}
|
|
|
|
Ok(Edit::range_replacement(
|
|
tree.codegen_stylist(stylist),
|
|
expr.range(),
|
|
))
|
|
}
|
|
|
|
/// (C414) Convert `sorted(list(foo))` to `sorted(foo)`
|
|
pub(crate) fn fix_unnecessary_double_cast_or_process(
|
|
expr: &Expr,
|
|
locator: &Locator,
|
|
stylist: &Stylist,
|
|
) -> Result<Edit> {
|
|
let module_text = locator.slice(expr);
|
|
let mut tree = match_expression(module_text)?;
|
|
let outer_call = match_call_mut(&mut tree)?;
|
|
|
|
outer_call.args = match outer_call.args.split_first() {
|
|
Some((first, rest)) => {
|
|
let inner_call = match_call(&first.value)?;
|
|
inner_call
|
|
.args
|
|
.iter()
|
|
.filter(|argument| argument.keyword.is_none())
|
|
.take(1)
|
|
.chain(rest.iter())
|
|
.cloned()
|
|
.collect::<Vec<_>>()
|
|
}
|
|
None => bail!("Expected at least one argument in outer function call"),
|
|
};
|
|
|
|
Ok(Edit::range_replacement(
|
|
tree.codegen_stylist(stylist),
|
|
expr.range(),
|
|
))
|
|
}
|
|
|
|
/// (C416) Convert `[i for i in x]` to `list(x)`.
|
|
pub(crate) fn fix_unnecessary_comprehension(
|
|
expr: &Expr,
|
|
locator: &Locator,
|
|
stylist: &Stylist,
|
|
) -> Result<Edit> {
|
|
let module_text = locator.slice(expr);
|
|
let mut tree = match_expression(module_text)?;
|
|
|
|
match &tree {
|
|
Expression::ListComp(inner) => {
|
|
tree = Expression::Call(Box::new(Call {
|
|
func: Box::new(Expression::Name(Box::new(Name {
|
|
value: "list",
|
|
lpar: vec![],
|
|
rpar: vec![],
|
|
}))),
|
|
args: vec![Arg {
|
|
value: inner.for_in.iter.clone(),
|
|
keyword: None,
|
|
equal: None,
|
|
comma: None,
|
|
star: "",
|
|
whitespace_after_star: ParenthesizableWhitespace::default(),
|
|
whitespace_after_arg: ParenthesizableWhitespace::default(),
|
|
}],
|
|
lpar: vec![],
|
|
rpar: vec![],
|
|
whitespace_after_func: ParenthesizableWhitespace::default(),
|
|
whitespace_before_args: ParenthesizableWhitespace::default(),
|
|
}));
|
|
}
|
|
Expression::SetComp(inner) => {
|
|
tree = Expression::Call(Box::new(Call {
|
|
func: Box::new(Expression::Name(Box::new(Name {
|
|
value: "set",
|
|
lpar: vec![],
|
|
rpar: vec![],
|
|
}))),
|
|
args: vec![Arg {
|
|
value: inner.for_in.iter.clone(),
|
|
keyword: None,
|
|
equal: None,
|
|
comma: None,
|
|
star: "",
|
|
whitespace_after_star: ParenthesizableWhitespace::default(),
|
|
whitespace_after_arg: ParenthesizableWhitespace::default(),
|
|
}],
|
|
lpar: vec![],
|
|
rpar: vec![],
|
|
whitespace_after_func: ParenthesizableWhitespace::default(),
|
|
whitespace_before_args: ParenthesizableWhitespace::default(),
|
|
}));
|
|
}
|
|
Expression::DictComp(inner) => {
|
|
tree = Expression::Call(Box::new(Call {
|
|
func: Box::new(Expression::Name(Box::new(Name {
|
|
value: "dict",
|
|
lpar: vec![],
|
|
rpar: vec![],
|
|
}))),
|
|
args: vec![Arg {
|
|
value: inner.for_in.iter.clone(),
|
|
keyword: None,
|
|
equal: None,
|
|
comma: None,
|
|
star: "",
|
|
whitespace_after_star: ParenthesizableWhitespace::default(),
|
|
whitespace_after_arg: ParenthesizableWhitespace::default(),
|
|
}],
|
|
lpar: vec![],
|
|
rpar: vec![],
|
|
whitespace_after_func: ParenthesizableWhitespace::default(),
|
|
whitespace_before_args: ParenthesizableWhitespace::default(),
|
|
}));
|
|
}
|
|
_ => {
|
|
bail!("Expected Expression::ListComp | Expression:SetComp | Expression:DictComp");
|
|
}
|
|
}
|
|
|
|
Ok(Edit::range_replacement(
|
|
pad(tree.codegen_stylist(stylist), expr.range(), locator),
|
|
expr.range(),
|
|
))
|
|
}
|
|
|
|
/// (C417) Convert `map(lambda x: x * 2, bar)` to `(x * 2 for x in bar)`.
|
|
pub(crate) fn fix_unnecessary_map(
|
|
expr: &Expr,
|
|
parent: Option<&Expr>,
|
|
object_type: ObjectType,
|
|
locator: &Locator,
|
|
stylist: &Stylist,
|
|
) -> Result<Edit> {
|
|
let module_text = locator.slice(expr);
|
|
let mut tree = match_expression(module_text)?;
|
|
let call = match_call_mut(&mut tree)?;
|
|
|
|
let (lambda, iter) = match call.args.as_slice() {
|
|
[call] => {
|
|
let call = match_call(&call.value)?;
|
|
let [lambda, iter] = call.args.as_slice() else {
|
|
bail!("Expected two arguments");
|
|
};
|
|
let lambda = match_lambda(&lambda.value)?;
|
|
let iter = &iter.value;
|
|
(lambda, iter)
|
|
}
|
|
[lambda, iter] => {
|
|
let lambda = match_lambda(&lambda.value)?;
|
|
let iter = &iter.value;
|
|
(lambda, iter)
|
|
}
|
|
_ => bail!("Expected a call or lambda"),
|
|
};
|
|
|
|
// Format the lambda target.
|
|
let target = match lambda.params.params.as_slice() {
|
|
// Ex) `lambda: x`
|
|
[] => AssignTargetExpression::Name(Box::new(Name {
|
|
value: "_",
|
|
lpar: vec![],
|
|
rpar: vec![],
|
|
})),
|
|
// Ex) `lambda x: y`
|
|
[param] => AssignTargetExpression::Name(Box::new(param.name.clone())),
|
|
// Ex) `lambda x, y: z`
|
|
params => AssignTargetExpression::Tuple(Box::new(Tuple {
|
|
elements: params
|
|
.iter()
|
|
.map(|param| Element::Simple {
|
|
value: Expression::Name(Box::new(param.name.clone())),
|
|
comma: None,
|
|
})
|
|
.collect(),
|
|
lpar: vec![],
|
|
rpar: vec![],
|
|
})),
|
|
};
|
|
|
|
// Parenthesize the iterator, if necessary, as in:
|
|
// ```python
|
|
// map(lambda x: x, y if y else z)
|
|
// ```
|
|
let iter = iter.clone();
|
|
let iter = if iter.lpar().is_empty()
|
|
&& iter.rpar().is_empty()
|
|
&& matches!(iter, Expression::IfExp(_) | Expression::Lambda(_))
|
|
{
|
|
iter.with_parens(LeftParen::default(), RightParen::default())
|
|
} else {
|
|
iter
|
|
};
|
|
|
|
let compfor = Box::new(CompFor {
|
|
target,
|
|
iter,
|
|
ifs: vec![],
|
|
inner_for_in: None,
|
|
asynchronous: None,
|
|
whitespace_before: space(),
|
|
whitespace_after_for: space(),
|
|
whitespace_before_in: space(),
|
|
whitespace_after_in: space(),
|
|
});
|
|
|
|
match object_type {
|
|
ObjectType::Generator => {
|
|
tree = Expression::GeneratorExp(Box::new(GeneratorExp {
|
|
elt: lambda.body.clone(),
|
|
for_in: compfor,
|
|
lpar: vec![LeftParen::default()],
|
|
rpar: vec![RightParen::default()],
|
|
}));
|
|
}
|
|
ObjectType::List => {
|
|
tree = Expression::ListComp(Box::new(ListComp {
|
|
elt: lambda.body.clone(),
|
|
for_in: compfor,
|
|
lbracket: LeftSquareBracket::default(),
|
|
rbracket: RightSquareBracket::default(),
|
|
lpar: vec![],
|
|
rpar: vec![],
|
|
}));
|
|
}
|
|
ObjectType::Set => {
|
|
tree = Expression::SetComp(Box::new(SetComp {
|
|
elt: lambda.body.clone(),
|
|
for_in: compfor,
|
|
lpar: vec![],
|
|
rpar: vec![],
|
|
lbrace: LeftCurlyBrace::default(),
|
|
rbrace: RightCurlyBrace::default(),
|
|
}));
|
|
}
|
|
ObjectType::Dict => {
|
|
let elements = match lambda.body.as_ref() {
|
|
Expression::Tuple(tuple) => &tuple.elements,
|
|
Expression::List(list) => &list.elements,
|
|
_ => {
|
|
bail!("Expected tuple or list for dictionary comprehension")
|
|
}
|
|
};
|
|
let [key, value] = elements.as_slice() else {
|
|
bail!("Expected container to include two elements");
|
|
};
|
|
let Element::Simple { value: key, .. } = key else {
|
|
bail!("Expected container to use a key as the first element");
|
|
};
|
|
let Element::Simple { value, .. } = value else {
|
|
bail!("Expected container to use a value as the second element");
|
|
};
|
|
|
|
tree = Expression::DictComp(Box::new(DictComp {
|
|
for_in: compfor,
|
|
lpar: vec![],
|
|
rpar: vec![],
|
|
key: Box::new(key.clone()),
|
|
value: Box::new(value.clone()),
|
|
lbrace: LeftCurlyBrace::default(),
|
|
rbrace: RightCurlyBrace::default(),
|
|
whitespace_before_colon: ParenthesizableWhitespace::default(),
|
|
whitespace_after_colon: ParenthesizableWhitespace::SimpleWhitespace(
|
|
SimpleWhitespace(" "),
|
|
),
|
|
}));
|
|
}
|
|
}
|
|
|
|
let mut content = tree.codegen_stylist(stylist);
|
|
|
|
// If the expression is embedded in an f-string, surround it with spaces to avoid
|
|
// syntax errors.
|
|
if matches!(object_type, ObjectType::Set | ObjectType::Dict) {
|
|
if parent.is_some_and(Expr::is_formatted_value_expr) {
|
|
content = format!(" {content} ");
|
|
}
|
|
}
|
|
|
|
Ok(Edit::range_replacement(content, expr.range()))
|
|
}
|
|
|
|
/// (C418) Convert `dict({"a": 1})` to `{"a": 1}`
|
|
pub(crate) fn fix_unnecessary_literal_within_dict_call(
|
|
expr: &Expr,
|
|
locator: &Locator,
|
|
stylist: &Stylist,
|
|
) -> Result<Edit> {
|
|
let module_text = locator.slice(expr);
|
|
let mut tree = match_expression(module_text)?;
|
|
let call = match_call_mut(&mut tree)?;
|
|
let arg = match_arg(call)?;
|
|
|
|
tree = arg.value.clone();
|
|
|
|
Ok(Edit::range_replacement(
|
|
tree.codegen_stylist(stylist),
|
|
expr.range(),
|
|
))
|
|
}
|
|
|
|
/// (C419) Convert `[i for i in a]` into `i for i in a`
|
|
pub(crate) fn fix_unnecessary_comprehension_any_all(
|
|
expr: &Expr,
|
|
locator: &Locator,
|
|
stylist: &Stylist,
|
|
) -> Result<Fix> {
|
|
// Expr(ListComp) -> Expr(GeneratorExp)
|
|
let module_text = locator.slice(expr);
|
|
let mut tree = match_expression(module_text)?;
|
|
let call = match_call_mut(&mut tree)?;
|
|
|
|
let (whitespace_after, whitespace_before, elt, for_in, lpar, rpar) = match &call.args[0].value {
|
|
Expression::ListComp(list_comp) => (
|
|
&list_comp.lbracket.whitespace_after,
|
|
&list_comp.rbracket.whitespace_before,
|
|
&list_comp.elt,
|
|
&list_comp.for_in,
|
|
&list_comp.lpar,
|
|
&list_comp.rpar,
|
|
),
|
|
Expression::SetComp(set_comp) => (
|
|
&set_comp.lbrace.whitespace_after,
|
|
&set_comp.rbrace.whitespace_before,
|
|
&set_comp.elt,
|
|
&set_comp.for_in,
|
|
&set_comp.lpar,
|
|
&set_comp.rpar,
|
|
),
|
|
_ => {
|
|
bail!("Expected Expression::ListComp | Expression::SetComp");
|
|
}
|
|
};
|
|
|
|
let mut new_empty_lines = vec![];
|
|
|
|
if let ParenthesizableWhitespace::ParenthesizedWhitespace(ParenthesizedWhitespace {
|
|
first_line,
|
|
empty_lines,
|
|
..
|
|
}) = &whitespace_after
|
|
{
|
|
// If there's a comment on the line after the opening bracket, we need
|
|
// to preserve it. The way we do this is by adding a new empty line
|
|
// with the same comment.
|
|
//
|
|
// Example:
|
|
// ```python
|
|
// any(
|
|
// [ # comment
|
|
// ...
|
|
// ]
|
|
// )
|
|
//
|
|
// # The above code will be converted to:
|
|
// any(
|
|
// # comment
|
|
// ...
|
|
// )
|
|
// ```
|
|
if let TrailingWhitespace {
|
|
comment: Some(comment),
|
|
..
|
|
} = first_line
|
|
{
|
|
// The indentation should be same as that of the opening bracket,
|
|
// but we don't have that information here. This will be addressed
|
|
// before adding these new nodes.
|
|
new_empty_lines.push(EmptyLine {
|
|
comment: Some(comment.clone()),
|
|
..EmptyLine::default()
|
|
});
|
|
}
|
|
if !empty_lines.is_empty() {
|
|
new_empty_lines.extend(empty_lines.clone());
|
|
}
|
|
}
|
|
|
|
if !new_empty_lines.is_empty() {
|
|
call.whitespace_before_args = match &call.whitespace_before_args {
|
|
ParenthesizableWhitespace::ParenthesizedWhitespace(ParenthesizedWhitespace {
|
|
first_line,
|
|
indent,
|
|
last_line,
|
|
..
|
|
}) => {
|
|
// Add the indentation of the opening bracket to all the new
|
|
// empty lines.
|
|
for empty_line in &mut new_empty_lines {
|
|
empty_line.whitespace = last_line.clone();
|
|
}
|
|
ParenthesizableWhitespace::ParenthesizedWhitespace(ParenthesizedWhitespace {
|
|
first_line: first_line.clone(),
|
|
empty_lines: new_empty_lines,
|
|
indent: *indent,
|
|
last_line: last_line.clone(),
|
|
})
|
|
}
|
|
// This is a rare case, but it can happen if the opening bracket
|
|
// is on the same line as the function call.
|
|
//
|
|
// Example:
|
|
// ```python
|
|
// any([
|
|
// ...
|
|
// ]
|
|
// )
|
|
// ```
|
|
ParenthesizableWhitespace::SimpleWhitespace(whitespace) => {
|
|
for empty_line in &mut new_empty_lines {
|
|
empty_line.whitespace = whitespace.clone();
|
|
}
|
|
ParenthesizableWhitespace::ParenthesizedWhitespace(ParenthesizedWhitespace {
|
|
empty_lines: new_empty_lines,
|
|
..ParenthesizedWhitespace::default()
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
let rbracket_comment =
|
|
if let ParenthesizableWhitespace::ParenthesizedWhitespace(ParenthesizedWhitespace {
|
|
first_line:
|
|
TrailingWhitespace {
|
|
whitespace,
|
|
comment: Some(comment),
|
|
..
|
|
},
|
|
..
|
|
}) = &whitespace_before
|
|
{
|
|
Some(format!("{}{}", whitespace.0, comment.0))
|
|
} else {
|
|
None
|
|
};
|
|
|
|
call.args[0].value = Expression::GeneratorExp(Box::new(GeneratorExp {
|
|
elt: elt.clone(),
|
|
for_in: for_in.clone(),
|
|
lpar: lpar.clone(),
|
|
rpar: rpar.clone(),
|
|
}));
|
|
|
|
let whitespace_after_arg = match &call.args[0].comma {
|
|
Some(comma) => {
|
|
let whitespace_after_comma = comma.whitespace_after.clone();
|
|
call.args[0].comma = None;
|
|
whitespace_after_comma
|
|
}
|
|
_ => call.args[0].whitespace_after_arg.clone(),
|
|
};
|
|
|
|
let new_comment;
|
|
call.args[0].whitespace_after_arg = match rbracket_comment {
|
|
Some(existing_comment) => {
|
|
if let ParenthesizableWhitespace::ParenthesizedWhitespace(ParenthesizedWhitespace {
|
|
first_line:
|
|
TrailingWhitespace {
|
|
whitespace: SimpleWhitespace(whitespace),
|
|
comment: Some(Comment(comment)),
|
|
..
|
|
},
|
|
empty_lines,
|
|
indent,
|
|
last_line,
|
|
}) = &whitespace_after_arg
|
|
{
|
|
new_comment = format!("{existing_comment}{whitespace}{comment}");
|
|
ParenthesizableWhitespace::ParenthesizedWhitespace(ParenthesizedWhitespace {
|
|
first_line: TrailingWhitespace {
|
|
comment: Some(Comment(new_comment.as_str())),
|
|
..TrailingWhitespace::default()
|
|
},
|
|
empty_lines: empty_lines.clone(),
|
|
indent: *indent,
|
|
last_line: last_line.clone(),
|
|
})
|
|
} else {
|
|
whitespace_after_arg
|
|
}
|
|
}
|
|
_ => whitespace_after_arg,
|
|
};
|
|
|
|
Ok(Fix::suggested(Edit::range_replacement(
|
|
tree.codegen_stylist(stylist),
|
|
expr.range(),
|
|
)))
|
|
}
|