rustpython_ruff_python_ast/
stmt_if.rs

1use std::iter;
2
3use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer};
4use ruff_text_size::{Ranged, TextRange};
5
6use crate::{ElifElseClause, Expr, Stmt, StmtIf};
7
8/// Return the `Range` of the first `Elif` or `Else` token in an `If` statement.
9pub fn elif_else_range(clause: &ElifElseClause, contents: &str) -> Option<TextRange> {
10    let token = SimpleTokenizer::new(contents, clause.range)
11        .skip_trivia()
12        .next()?;
13    matches!(token.kind, SimpleTokenKind::Elif | SimpleTokenKind::Else).then_some(token.range())
14}
15
16#[derive(Copy, Clone, Debug, PartialEq, Eq)]
17pub enum BranchKind {
18    If,
19    Elif,
20}
21
22#[derive(Debug)]
23pub struct IfElifBranch<'a> {
24    pub kind: BranchKind,
25    pub test: &'a Expr,
26    pub body: &'a [Stmt],
27    range: TextRange,
28}
29
30impl Ranged for IfElifBranch<'_> {
31    fn range(&self) -> TextRange {
32        self.range
33    }
34}
35
36pub fn if_elif_branches(stmt_if: &StmtIf) -> impl Iterator<Item = IfElifBranch<'_>> {
37    iter::once(IfElifBranch {
38        kind: BranchKind::If,
39        test: stmt_if.test.as_ref(),
40        body: stmt_if.body.as_slice(),
41        range: TextRange::new(stmt_if.start(), stmt_if.body.last().unwrap().end()),
42    })
43    .chain(stmt_if.elif_else_clauses.iter().filter_map(|clause| {
44        Some(IfElifBranch {
45            kind: BranchKind::Elif,
46            test: clause.test.as_ref()?,
47            body: clause.body.as_slice(),
48            range: clause.range,
49        })
50    }))
51}
rustpython_ruff_python_ast/stmt_if.rs

rustpython_ruff_python_ast/
stmt_if.rs
