Repository: astral-sh/ruff
File: crates/ruff_linter/src/rules/pylint/rules/stop_iteration_return.rs (Rust)

Code before:
            ast::Stmt::Raise(stmt_raise) => {
                if is_stop_iteration_raise(stmt_raise, self.checker) {
                    self.stop_iteration_raises.push(stmt_raise);
                }
                walk_stmt(self, stmt);
            }
            _ => walk_stmt(self, stmt),
        }
    }

    fn visit_expr(&mut self, expr: &'a ast::Expr) {
        match expr {
            ast::Expr::Lambda(_) => {}
            ast::Expr::Yield(_) | ast::Expr::YieldFrom(_) => {
                self.has_yield = true;
                walk_expr(self, expr);
            }
            _ => walk_expr(self, expr),
        }
    }
}

fn is_stop_iteration_raise(stmt_raise: &ast::StmtRaise, checker: &Checker) -> bool {
    let Some(exc) = &stmt_raise.exc else {
        return false;
    };

    match exc.as_ref() {
        ast::Expr::Call(ast::ExprCall { func, .. }) => {
            checker.semantic().match_builtin_expr(func, "StopIteration")
        }
        expr => checker.semantic().match_builtin_expr(expr, "StopIteration"),
    }
}


Review comment (bug):
I'm not sure how helpful this helper function is. I think we could just destructure the `StmtRaise` in the `match` on line 83 and use `map_callable`:

https://github.com/astral-sh/ruff/blob/192c37d5401786b3043aafbfb657b4f5e7df8ac8/crates/ruff_python_ast/src/helpers.rs#L724-L735

For example,

```rust
    fn visit_stmt(&mut self, stmt: &'a ast::Stmt) {
        match stmt {
            ast::Stmt::FunctionDef(_) => {}
            ast::Stmt::Raise(ast::StmtRaise { exc: Some(exc), .. }) => {
                if self.checker.semantic().match_builtin_expr(map_callable(exc), "StopIteration") {
                    self.stop_iteration_raises.push(stmt_raise);
                }
                walk_stmt(self, stmt);
            }
            _ => walk_stmt(self, stmt),
        }
    }
```

I was going to suggest removing the `walk_stmt` call in the `Stmt::Raise` arm, but I think that's actually correct, and B901 has a bug in it by only checking for statements with a single yield expression:

https://github.com/astral-sh/ruff/blob/192c37d5401786b3043aafbfb657b4f5e7df8ac8/crates/ruff_linter/src/rules/flake8_bugbear/rules/return_in_generator.rs#L118-L123

That means it misses cases like this:

```py
def foo():
    return (yield 1)
```

Feel free to follow up on that if you're interested!

This also applies here, where your code will correctly flag something like this:

```py
def foo():
    raise StopIteration((yield 1))
```

That might be a nice test case to add.

Code after:
use ruff_macros::{ViolationMetadata, derive_message_formats};
use ruff_python_ast::{
    self as ast,
    helpers::map_callable,
    visitor::{Visitor, walk_expr, walk_stmt},
};
use ruff_text_size::Ranged;

use crate::Violation;
use crate::checkers::ast::Checker;

/// ## What it does
/// Checks for explicit `raise StopIteration` in generator functions.
///
/// ## Why is this bad?
/// Raising `StopIteration` in a generator function causes a `RuntimeError`
/// when the generator is iterated over.
///
/// Instead of `raise StopIteration`, use `return` in generator functions.
///
/// ## Example
/// ```python
/// def my_generator():
///     yield 1
///     yield 2
///     raise StopIteration  # This causes RuntimeError at runtime
/// ```
///
/// Use instead:
/// ```python
/// def my_generator():
///     yield 1
///     yield 2
///     return  # Use return instead
/// ```
///
/// ## References
/// - [PEP 479](https://peps.python.org/pep-0479/)
/// - [Python documentation](https://docs.python.org/3/library/exceptions.html#StopIteration)
#[derive(ViolationMetadata)]
#[violation_metadata(preview_since = "0.14.3")]
pub(crate) struct StopIterationReturn;

impl Violation for StopIterationReturn {
    #[derive_message_formats]
    fn message(&self) -> String {
        "Explicit `raise StopIteration` in generator".to_string()
    }

    fn fix_title(&self) -> Option<String> {
        Some("Use `return` instead".to_string())
    }
}

/// PLR1708
pub(crate) fn stop_iteration_return(checker: &Checker, function_def: &ast::StmtFunctionDef) {
    let mut analyzer = GeneratorAnalyzer {
        checker,
        has_yield: false,
        stop_iteration_raises: Vec::new(),
    };

    analyzer.visit_body(&function_def.body);

    if analyzer.has_yield {
        for raise_stmt in analyzer.stop_iteration_raises {
            checker.report_diagnostic(StopIterationReturn, raise_stmt.range());
        }
    }
}

struct GeneratorAnalyzer<'a, 'b> {
    checker: &'a Checker<'b>,
    has_yield: bool,
    stop_iteration_raises: Vec<&'a ast::StmtRaise>,
}

impl<'a> Visitor<'a> for GeneratorAnalyzer<'a, '_> {
    fn visit_stmt(&mut self, stmt: &'a ast::Stmt) {
        match stmt {
            ast::Stmt::FunctionDef(_) => {}
            ast::Stmt::Raise(raise @ ast::StmtRaise { exc: Some(exc), .. }) => {
                if self
                    .checker
                    .semantic()
                    .match_builtin_expr(map_callable(exc), "StopIteration")
                {
                    self.stop_iteration_raises.push(raise);
                }
                walk_stmt(self, stmt);
            }
            _ => walk_stmt(self, stmt),
        }
    }

    fn visit_expr(&mut self, expr: &'a ast::Expr) {
        match expr {
            ast::Expr::Lambda(_) => {}
            ast::Expr::Yield(_) | ast::Expr::YieldFrom(_) => {
                self.has_yield = true;
                walk_expr(self, expr);
            }
            _ => walk_expr(self, expr),
        }
    }
}