diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_simplify/SIM113.py b/crates/ruff_linter/resources/test/fixtures/flake8_simplify/SIM113.py index e0ac4190ed..00771b5f70 100644 --- a/crates/ruff_linter/resources/test/fixtures/flake8_simplify/SIM113.py +++ b/crates/ruff_linter/resources/test/fixtures/flake8_simplify/SIM113.py @@ -46,7 +46,8 @@ def func(): def func(): - # OK (index doesn't start at 0 + # SIM113 + # https://github.com/astral-sh/ruff/pull/21395 idx = 10 for x in range(5): g(x, idx) diff --git a/crates/ruff_linter/src/preview.rs b/crates/ruff_linter/src/preview.rs index 836ba4feea..239466f599 100644 --- a/crates/ruff_linter/src/preview.rs +++ b/crates/ruff_linter/src/preview.rs @@ -269,3 +269,8 @@ pub(crate) const fn is_typing_extensions_str_alias_enabled(settings: &LinterSett pub(crate) const fn is_extended_i18n_function_matching_enabled(settings: &LinterSettings) -> bool { settings.preview.is_enabled() } + +// https://github.com/astral-sh/ruff/pull/21395 +pub(crate) const fn is_enumerate_for_loop_int_index_enabled(settings: &LinterSettings) -> bool { + settings.preview.is_enabled() +} diff --git a/crates/ruff_linter/src/rules/flake8_simplify/mod.rs b/crates/ruff_linter/src/rules/flake8_simplify/mod.rs index 45233277e5..4546dd143a 100644 --- a/crates/ruff_linter/src/rules/flake8_simplify/mod.rs +++ b/crates/ruff_linter/src/rules/flake8_simplify/mod.rs @@ -61,6 +61,7 @@ mod tests { #[test_case(Rule::SplitStaticString, Path::new("SIM905.py"))] #[test_case(Rule::DictGetWithNoneDefault, Path::new("SIM910.py"))] + #[test_case(Rule::EnumerateForLoop, Path::new("SIM113.py"))] fn preview_rules(rule_code: Rule, path: &Path) -> Result<()> { let snapshot = format!( "preview__{}_{}", diff --git a/crates/ruff_linter/src/rules/flake8_simplify/rules/enumerate_for_loop.rs b/crates/ruff_linter/src/rules/flake8_simplify/rules/enumerate_for_loop.rs index 6739fa3868..a35513de85 100644 --- a/crates/ruff_linter/src/rules/flake8_simplify/rules/enumerate_for_loop.rs +++ b/crates/ruff_linter/src/rules/flake8_simplify/rules/enumerate_for_loop.rs @@ -1,6 +1,8 @@ +use crate::preview::is_enumerate_for_loop_int_index_enabled; use ruff_macros::{ViolationMetadata, derive_message_formats}; use ruff_python_ast::statement_visitor::{StatementVisitor, walk_stmt}; use ruff_python_ast::{self as ast, Expr, Int, Number, Operator, Stmt}; +use ruff_python_semantic::analyze::type_inference::{NumberLike, PythonType, ResolvedPythonType}; use ruff_python_semantic::analyze::typing; use ruff_text_size::Ranged; @@ -11,6 +13,9 @@ use crate::checkers::ast::Checker; /// Checks for `for` loops with explicit loop-index variables that can be replaced /// with `enumerate()`. /// +/// In [preview], this rule checks for index variables initialized with any integer rather than only +/// a literal zero. +/// /// ## Why is this bad? /// When iterating over a sequence, it's often desirable to keep track of the /// index of each element alongside the element itself. Prefer the `enumerate` @@ -35,6 +40,8 @@ use crate::checkers::ast::Checker; /// /// ## References /// - [Python documentation: `enumerate`](https://docs.python.org/3/library/functions.html#enumerate) +/// +/// [preview]: https://docs.astral.sh/ruff/preview/ #[derive(ViolationMetadata)] #[violation_metadata(stable_since = "v0.2.0")] pub(crate) struct EnumerateForLoop { @@ -82,17 +89,21 @@ pub(crate) fn enumerate_for_loop(checker: &Checker, for_stmt: &ast::StmtFor) { continue; } - // Ensure that the index variable was initialized to 0. + // Ensure that the index variable was initialized to 0 (or instance of `int` if preview is enabled). let Some(value) = typing::find_binding_value(binding, checker.semantic()) else { continue; }; - if !matches!( + if !(matches!( value, Expr::NumberLiteral(ast::ExprNumberLiteral { value: Number::Int(Int::ZERO), .. }) - ) { + ) || matches!( + ResolvedPythonType::from(value), + ResolvedPythonType::Atom(PythonType::Number(NumberLike::Integer)) + ) && is_enumerate_for_loop_int_index_enabled(checker.settings())) + { continue; } diff --git a/crates/ruff_linter/src/rules/flake8_simplify/snapshots/ruff_linter__rules__flake8_simplify__tests__preview__SIM113_SIM113.py.snap b/crates/ruff_linter/src/rules/flake8_simplify/snapshots/ruff_linter__rules__flake8_simplify__tests__preview__SIM113_SIM113.py.snap new file mode 100644 index 0000000000..065ed20bb9 --- /dev/null +++ b/crates/ruff_linter/src/rules/flake8_simplify/snapshots/ruff_linter__rules__flake8_simplify__tests__preview__SIM113_SIM113.py.snap @@ -0,0 +1,60 @@ +--- +source: crates/ruff_linter/src/rules/flake8_simplify/mod.rs +--- +SIM113 Use `enumerate()` for index variable `idx` in `for` loop + --> SIM113.py:6:9 + | +4 | for x in range(5): +5 | g(x, idx) +6 | idx += 1 + | ^^^^^^^^ +7 | h(x) + | + +SIM113 Use `enumerate()` for index variable `idx` in `for` loop + --> SIM113.py:17:9 + | +15 | if g(x): +16 | break +17 | idx += 1 + | ^^^^^^^^ +18 | sum += h(x, idx) + | + +SIM113 Use `enumerate()` for index variable `idx` in `for` loop + --> SIM113.py:27:9 + | +25 | g(x) +26 | h(x, y) +27 | idx += 1 + | ^^^^^^^^ + | + +SIM113 Use `enumerate()` for index variable `idx` in `for` loop + --> SIM113.py:36:9 + | +34 | for x in range(5): +35 | sum += h(x, idx) +36 | idx += 1 + | ^^^^^^^^ + | + +SIM113 Use `enumerate()` for index variable `idx` in `for` loop + --> SIM113.py:44:9 + | +42 | for x in range(5): +43 | g(x, idx) +44 | idx += 1 + | ^^^^^^^^ +45 | h(x) + | + +SIM113 Use `enumerate()` for index variable `idx` in `for` loop + --> SIM113.py:54:9 + | +52 | for x in range(5): +53 | g(x, idx) +54 | idx += 1 + | ^^^^^^^^ +55 | h(x) + |