Skip to content

Commit

Permalink
Allow fix with start
Browse files Browse the repository at this point in the history
  • Loading branch information
@charliermarsh
charliermarsh committed Sep 19, 2023
1 parent 1178a4b commit faa4576
Show file tree
Hide file tree
Showing 3 changed files with 213 additions and 155 deletions.
12 changes: 12 additions & 0 deletions crates/ruff/resources/test/fixtures/refurb/FURB148.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,12 @@
for index, _ in enumerate(books, 1):
print(index)

for index, _ in enumerate(books, start=x):
print(book)

for index, _ in enumerate(books, x):
print(book)

for _, book in enumerate(books):
print(book)

Expand All @@ -31,6 +37,12 @@
for _, book in enumerate(books, 1):
print(book)

for _, book in enumerate(books, start=x):
print(book)

for _, book in enumerate(books, x):
print(book)

for index, (_, _) in enumerate(books):
print(index)

Expand Down
100 changes: 43 additions & 57 deletions crates/ruff/src/rules/refurb/rules/unnecessary_enumerate.rs
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use std::fmt;

use num_traits::ToPrimitive;
use num_traits::Zero;

use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Fix, Violation};
use ruff_macros::{derive_message_formats, violation};
Expand Down Expand Up @@ -130,20 +130,16 @@ pub(crate) fn unnecessary_enumerate(checker: &mut Checker, stmt_for: &ast::StmtF

// The index is unused, so replace with `for value in sequence`.
if checker.patch(diagnostic.kind.rule()) {
// Get the `start` argument, if it is a constant integer.
if start(arguments).map_or(true, |start| start == 0) {
let replace_iter =
Edit::range_replacement(sequence.into(), stmt_for.iter.range());
let replace_target = Edit::range_replacement(
pad(
checker.locator().slice(value).to_string(),
stmt_for.target.range(),
checker.locator(),
),
let replace_iter = Edit::range_replacement(sequence.into(), stmt_for.iter.range());
let replace_target = Edit::range_replacement(
pad(
checker.locator().slice(value).to_string(),
stmt_for.target.range(),
);
diagnostic.set_fix(Fix::suggested_edits(replace_iter, [replace_target]));
}
checker.locator(),
),
stmt_for.target.range(),
);
diagnostic.set_fix(Fix::suggested_edits(replace_iter, [replace_target]));
}

checker.diagnostics.push(diagnostic);
Expand All @@ -156,25 +152,40 @@ pub(crate) fn unnecessary_enumerate(checker: &mut Checker, stmt_for: &ast::StmtF
},
func.range(),
);
if checker.patch(diagnostic.kind.rule()) {
// Get the `start` argument, if it is a constant integer.
let start = start(arguments);

let replace_iter = Edit::range_replacement(
generate_range_len_call(sequence, start, checker.generator()),
stmt_for.iter.range(),
);
if checker.patch(diagnostic.kind.rule())
&& checker.semantic().is_builtin("range")
&& checker.semantic().is_builtin("len")
{
// If the `start` argument is set to something other than the `range` default,
// there's no clear fix.
let start = arguments.find_argument("start", 1);
if start.map_or(true, |start| {
if let Expr::Constant(ast::ExprConstant {
value: Constant::Int(value),
..
}) = start
{
value.is_zero()
} else {
false
}
}) {
let replace_iter = Edit::range_replacement(
generate_range_len_call(sequence, checker.generator()),
stmt_for.iter.range(),
);

let replace_target = Edit::range_replacement(
pad(
checker.locator().slice(index).to_string(),
let replace_target = Edit::range_replacement(
pad(
checker.locator().slice(index).to_string(),
stmt_for.target.range(),
checker.locator(),
),
stmt_for.target.range(),
checker.locator(),
),
stmt_for.target.range(),
);
);

diagnostic.set_fix(Fix::suggested_edits(replace_iter, [replace_target]));
diagnostic.set_fix(Fix::suggested_edits(replace_iter, [replace_target]));
}
}
checker.diagnostics.push(diagnostic);
}
Expand All @@ -198,24 +209,9 @@ impl fmt::Display for EnumerateSubset {
}
}

/// Returns the value of the `start` argument to `enumerate`, if it is a
/// constant integer. Otherwise, returns `None`.
fn start(arguments: &Arguments) -> Option<u32> {
let step_param = arguments.find_argument("start", 1)?;
if let Expr::Constant(ast::ExprConstant {
value: Constant::Int(value),
..
}) = &step_param
{
value.to_u32()
} else {
None
}
}

/// Format a code snippet to call `range(len(name))`, where `name` is the given
/// sequence name.
fn generate_range_len_call(name: &str, start: Option<u32>, generator: Generator) -> String {
fn generate_range_len_call(name: &str, generator: Generator) -> String {
// Construct `name`.
let var = ast::ExprName {
id: name.to_string(),
Expand All @@ -240,16 +236,6 @@ fn generate_range_len_call(name: &str, start: Option<u32>, generator: Generator)
range: TextRange::default(),
};
// Construct `range(len(name))`.
let range_args: Vec<Expr> = match start {
Some(start) if start > 0 => {
let start_expr: Expr = Expr::Constant(ast::ExprConstant {
range: TextRange::default(),
value: Constant::Int(start.into()),
});
vec![start_expr, len.into()]
}
_ => vec![len.into()],
};
let range = ast::ExprCall {
func: Box::new(
ast::ExprName {
Expand All @@ -260,7 +246,7 @@ fn generate_range_len_call(name: &str, start: Option<u32>, generator: Generator)
.into(),
),
arguments: Arguments {
args: range_args,
args: vec![len.into()],
keywords: vec![],
range: TextRange::default(),
},
Expand Down

0 comments on commit faa4576

Please sign in to comment.