[flake8-simplify] Implement enumerate-for-loop (SIM113)

crates/ruff_linter/resources/test/fixtures/flake8_simplify/SIM113.py

@@ -0,0 +1,104 @@
+
+def f():
+    # SIM113
+    idx = 0
+    for x in iterable:
+        g(x, idx)
+        idx +=1
+        h(x)
+
+
+def f():
+    # SIM113
+    sum = 0
+    idx = 0
+    for x in iterable:
+        if g(x):
+            break
+        idx += 1
+        sum += h(x, idx)
+
+
+def f():
+    # SIM113
+    idx = 0
+    for x, y in iterable_tuples():
+        g(x)
+        h(x, y)
+        idx += 1
+
+def f():
+    # Current SIM113 doesn't catch this yet because for loop
+    # doesn't immidiately follow index initialization
+    idx = 0
+    sum = 0
+    for x in iterable:
+        sum += h(x, idx)
+        idx += 1
+
+
+def f():
+    # Current SIM113 doesn't catch this due to unpacking in
+    # in intialization
+    sum, idx = 0, 0
+    for x in iterable:
+        g(x, idx)
+        idx +=1
+        h(x)
+
+
+def f():
+    # loop doesn't start at zero
+    idx = 10
+    for x in iterable:
+        g(x, idx)
+        idx +=1
+        h(x)
+
+def f():
+    # index doesn't increment by one
+    idx = 0
+    for x in iterable:
+        g(x, idx)
+        idx +=2
+        h(x)
+
+def f():
+    # index increment inside condition
+    idx = 0
+    for x in iterable:
+        if g(x):
+            idx += 1
+        h(x)
+
+def f():
+    # Continue in match-case
+    idx = 0
+    for x in iterable:
+        match g(x):
+            case 1: h(x)
+            case 2: continue
+            case _: h(idx)
+        idx += 1
+
+def f():
+    # Continue inside with clause
+    idx = 0
+    for x in iterable:
+        with context as c:
+            if g(x):
+                continue
+            h(x, idx, c)
+        idx += 1
+
+
+def f():
+    # Continue in try block
+    idx = 0
+    for x in iterable:
+        try:
+            g(x, idx)
+        except:
+            h(x)
+            continue
+        idx += 1

crates/ruff_linter/src/checkers/ast/analyze/statement.rs

@@ -1236,6 +1236,9 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                 if checker.enabled(Rule::TryExceptInLoop) {
                     perflint::rules::try_except_in_loop(checker, body);
                 }
+                if checker.enabled(Rule::EnumerateForLoop) {
+                    flake8_simplify::rules::use_enumerate_in_for_loop(checker, stmt);
+                }
             }
         }
         Stmt::Try(ast::StmtTry {

crates/ruff_linter/src/codes.rs

@@ -423,6 +423,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Flake8Simplify, "109") => (RuleGroup::Unspecified, rules::flake8_simplify::rules::CompareWithTuple),
         (Flake8Simplify, "110") => (RuleGroup::Unspecified, rules::flake8_simplify::rules::ReimplementedBuiltin),
         (Flake8Simplify, "112") => (RuleGroup::Unspecified, rules::flake8_simplify::rules::UncapitalizedEnvironmentVariables),
+        (Flake8Simplify, "113") => (RuleGroup::Unspecified, rules::flake8_simplify::rules::EnumerateForLoop),
         (Flake8Simplify, "114") => (RuleGroup::Unspecified, rules::flake8_simplify::rules::IfWithSameArms),
         (Flake8Simplify, "115") => (RuleGroup::Unspecified, rules::flake8_simplify::rules::OpenFileWithContextHandler),
         (Flake8Simplify, "116") => (RuleGroup::Unspecified, rules::flake8_simplify::rules::IfElseBlockInsteadOfDictLookup),

crates/ruff_linter/src/rules/flake8_simplify/mod.rs

@@ -26,6 +26,7 @@ mod tests {
     #[test_case(Rule::ReimplementedBuiltin, Path::new("SIM110.py"))]
     #[test_case(Rule::ReimplementedBuiltin, Path::new("SIM111.py"))]
     #[test_case(Rule::UncapitalizedEnvironmentVariables, Path::new("SIM112.py"))]
+    #[test_case(Rule::EnumerateForLoop, Path::new("SIM113.py"))]
     #[test_case(Rule::OpenFileWithContextHandler, Path::new("SIM115.py"))]
     #[test_case(Rule::MultipleWithStatements, Path::new("SIM117.py"))]
     #[test_case(Rule::InDictKeys, Path::new("SIM118.py"))]

crates/ruff_linter/src/rules/flake8_simplify/rules/ast_bool_op.rs

@@ -288,7 +288,7 @@ impl AlwaysFixableViolation for ExprAndFalse {
 }
 
 /// Return `true` if two `Expr` instances are equivalent names.
-fn is_same_expr<'a>(a: &'a Expr, b: &'a Expr) -> Option<&'a str> {
+pub(crate) fn is_same_expr<'a>(a: &'a Expr, b: &'a Expr) -> Option<&'a str> {
     if let (Expr::Name(ast::ExprName { id: a, .. }), Expr::Name(ast::ExprName { id: b, .. })) =
         (&a, &b)
     {

crates/ruff_linter/src/rules/flake8_simplify/rules/enumerate_for.rs

@@ -0,0 +1,178 @@
+use ruff_python_ast::{self as ast, Constant, ExceptHandler, Expr, Int, MatchCase, Operator, Stmt};
+use ruff_text_size::{Ranged, TextRange};
+
+use crate::rules::flake8_simplify::rules::ast_bool_op::is_same_expr;
+use ruff_diagnostics::{Diagnostic, Violation};
+use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::traversal;
+
+use crate::checkers::ast::Checker;
+
+/// ## What it does
+/// Checks for loops which has explicit loop-index variables that can be simplified by using `enumerate()`.
+///
+/// ## Why this is bad?
+/// Using `enumerate()` is more readable and concise. It could lead to more efficient
+/// and less error-prone code.
+///
+/// ## Example
+/// ```python
+/// sum = 0
+/// idx = 0
+/// for item in items:
+///    sum += func(item, idx)
+///   idx += 1
+/// ```
+///
+/// Use instead:
+/// ```python
+/// sum = 0
+/// for idx, item in enumerate(items):
+///    sum += func(item, idx)
+/// ```
+///
+/// ## References
+/// - [Python documentation: enumerate()](https://docs.python.org/3/library/functions.html#enumerate)
+#[violation]
+pub struct EnumerateForLoop {
+    index_var: String,
+}
+
+impl Violation for EnumerateForLoop {
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        let EnumerateForLoop { index_var } = self;
+        format!("Use enumereate() for index variable `{index_var}` in for loop")
+    }
+}
+
+/// SIM113
+pub(crate) fn use_enumerate_in_for_loop(checker: &mut Checker, stmt: &Stmt) {
+    if !checker.semantic().current_scope().kind.is_function() {
+        return;
+    }
+    let Stmt::For(ast::StmtFor { body, .. }) = stmt else {
+        return;
+    };
+
+    // Check if loop body contains a continue statement
+    if loop_possibly_continues(body) {
+        return;
+    };
+    // Check if index variable is initialized to zero prior to loop
+    let Some((prev_stmt, index_var)) = get_candidate_loop_index(checker, stmt) else {
+        return;
+    };
+
+    // Check if loop body contains an index increment statement matching `index_var
+    if body.iter().any(|stmt| is_index_increment(stmt, index_var)) {
+        let diagonastic = Diagnostic::new(
+            EnumerateForLoop {
+                index_var: checker.generator().expr(index_var),
+            },
+            TextRange::new(prev_stmt.start(), stmt.end()),
+        );
+        checker.diagnostics.push(diagonastic);
+    }
+}
+
+/// Recursively check if the `for` loop body contains a `continue` statement
+fn loop_possibly_continues(body: &[Stmt]) -> bool {
+    body.iter().any(|stmt| match stmt {
+        Stmt::Continue(_) => true,
+        Stmt::If(ast::StmtIf {
+            body,
+            elif_else_clauses,
+            ..
+        }) => {
+            loop_possibly_continues(body)
+                || elif_else_clauses
+                    .iter()
+                    .any(|clause| loop_possibly_continues(&clause.body))
+        }
+        Stmt::With(ast::StmtWith { body, .. }) => loop_possibly_continues(body),
+        Stmt::Match(ast::StmtMatch { cases, .. }) => cases
+            .iter()
+            .any(|MatchCase { body, .. }| loop_possibly_continues(body)),
+        Stmt::Try(ast::StmtTry {
+            body,
+            handlers,
+            orelse,
+            finalbody,
+            ..
+        }) => {
+            loop_possibly_continues(body)
+                || loop_possibly_continues(orelse)
+                || loop_possibly_continues(finalbody)
+                || handlers.iter().any(|handler| match handler {
+                    ExceptHandler::ExceptHandler(ast::ExceptHandlerExceptHandler {
+                        body, ..
+                    }) => loop_possibly_continues(body),
+                })
+        }
+
+        _ => false,
+    })
+}
+
+/// Check previous statement of `for` loop to find a possible index variable
+/// which is initialized to zero
+/// Ex:
+/// ```python
+/// idx = 0
+/// for item in items:
+///     ...
+/// ```
+/// Return (&Stmt, &Expr) to initialization stmt and `idx` variable
+fn get_candidate_loop_index<'a>(
+    checker: &'a Checker,
+    stmt: &'a Stmt,
+) -> Option<(&'a Stmt, &'a Expr)> {
+    let parent = checker.semantic().current_statement_parent()?;
+    let suite = traversal::suite(stmt, parent)?;
+    let prev_stmt = traversal::prev_sibling(stmt, suite)?;
+    // check if it's a possible index initialization i.e. idx = 0
+    let Stmt::Assign(ast::StmtAssign { targets, value, .. }) = prev_stmt else {
+        return None;
+    };
+    let [Expr::Name(ast::ExprName { id: _, .. })] = targets.as_slice() else {
+        return None;
+    };
+    if let Expr::Constant(ast::ExprConstant {
+        value: Constant::Int(value),
+        ..
+    }) = value.as_ref()
+    {
+        if matches!(*value, Int::ZERO) {
+            return Some((prev_stmt, &targets[0]));
+        }
+    }
+
+    None
+}
+
+// Check if `stmt` is `index_var` += 1
+fn is_index_increment(stmt: &Stmt, index_var: &Expr) -> bool {
+    let Stmt::AugAssign(ast::StmtAugAssign {
+        target, op, value, ..
+    }) = stmt
+    else {
+        return false;
+    };
+    if !matches!(op, Operator::Add) {
+        return false;
+    }
+    let Some(_) = is_same_expr(index_var, target) else {
+        return false;
+    };
+    if let Expr::Constant(ast::ExprConstant {
+        value: Constant::Int(value),
+        ..
+    }) = value.as_ref()
+    {
+        if matches!(*value, Int::ONE) {
+            return true;
+        }
+    }
+    false
+}

crates/ruff_linter/src/rules/flake8_simplify/rules/mod.rs

@@ -4,6 +4,7 @@ pub(crate) use ast_ifexp::*;
 pub(crate) use ast_unary_op::*;
 pub(crate) use ast_with::*;
 pub(crate) use collapsible_if::*;
+pub(crate) use enumerate_for::*;
 pub(crate) use if_else_block_instead_of_dict_get::*;
 pub(crate) use if_else_block_instead_of_dict_lookup::*;
 pub(crate) use if_else_block_instead_of_if_exp::*;
@@ -22,6 +23,7 @@ mod ast_ifexp;
 mod ast_unary_op;
 mod ast_with;
 mod collapsible_if;
+mod enumerate_for;
 mod fix_with;
 mod if_else_block_instead_of_dict_get;
 mod if_else_block_instead_of_dict_lookup;

crates/ruff_linter/src/rules/flake8_simplify/snapshots/ruff_linter__rules__flake8_simplify__tests__SIM113_SIM113.py.snap

@@ -0,0 +1,46 @@
+---
+source: crates/ruff_linter/src/rules/flake8_simplify/mod.rs
+---
+SIM113.py:4:5: SIM113 Use enumereate() for index variable `idx` in for loop
+  |
+2 |   def f():
+3 |       # SIM113
+4 |       idx = 0
+  |  _____^
+5 | |     for x in iterable:
+6 | |         g(x, idx)
+7 | |         idx +=1
+8 | |         h(x)
+  | |____________^ SIM113
+  |
+
+SIM113.py:14:5: SIM113 Use enumereate() for index variable `idx` in for loop
+   |
+12 |       # SIM113
+13 |       sum = 0
+14 |       idx = 0
+   |  _____^
+15 | |     for x in iterable:
+16 | |         if g(x):
+17 | |             break
+18 | |         idx += 1
+19 | |         sum += h(x, idx)
+   | |________________________^ SIM113
+   |
+
+SIM113.py:24:5: SIM113 Use enumereate() for index variable `idx` in for loop
+   |
+22 |   def f():
+23 |       # SIM113
+24 |       idx = 0
+   |  _____^
+25 | |     for x, y in iterable_tuples():
+26 | |         g(x)
+27 | |         h(x, y)
+28 | |         idx += 1
+   | |________________^ SIM113
+29 |           
+30 |   def f():
+   |
+
+

crates/ruff_python_ast/src/traversal.rs

@@ -79,3 +79,15 @@ pub fn next_sibling<'a>(stmt: &'a Stmt, suite: &'a Suite) -> Option<&'a Stmt> {
     }
     None
 }
+
+/// Given a [`Stmt`] and its containing [`Suite`], return the previous [`Stmt`] in the [`Suite`].
+pub fn prev_sibling<'a>(stmt: &'a Stmt, suite: &'a Suite) -> Option<&'a Stmt> {
+    let mut prev = None;
+    for sibling in suite {
+        if sibling == stmt {
+            return prev;
+        }
+        prev = Some(sibling);
+    }
+    None
+}