Refactor StmtIf: Formatter and Linter (#5459)

## Summary

Previously, `StmtIf` was defined recursively as
```rust
pub struct StmtIf {
    pub range: TextRange,
    pub test: Box<Expr>,
    pub body: Vec<Stmt>,
    pub orelse: Vec<Stmt>,
}
```
Every `elif` was represented as an `orelse` with a single `StmtIf`. This
means that this representation couldn't differentiate between
```python
if cond1:
    x = 1
else:
    if cond2:
        x = 2
```
and 
```python
if cond1:
    x = 1
elif cond2:
    x = 2
```
It also makes many checks harder than they need to be because we have to
recurse just to iterate over an entire if-elif-else and because we're
lacking nodes and ranges on the `elif` and `else` branches.

We change the representation to a flat

```rust
pub struct StmtIf {
    pub range: TextRange,
    pub test: Box<Expr>,
    pub body: Vec<Stmt>,
    pub elif_else_clauses: Vec<ElifElseClause>,
}

pub struct ElifElseClause {
    pub range: TextRange,
    pub test: Option<Expr>,
    pub body: Vec<Stmt>,
}
```
where `test: Some(_)` represents an `elif` and `test: None` an else.

This representation is different tradeoff, e.g. we need to allocate the
`Vec<ElifElseClause>`, the `elif`s are now different than the `if`s
(which matters in rules where want to check both `if`s and `elif`s) and
the type system doesn't guarantee that the `test: None` else is actually
last. We're also now a bit more inconsistent since all other `else`,
those from `for`, `while` and `try`, still don't have nodes. With the
new representation some things became easier, e.g. finding the `elif`
token (we can use the start of the `ElifElseClause`) and formatting
comments for if-elif-else (no more dangling comments splitting, we only
have to insert the dangling comment after the colon manually and set
`leading_alternate_branch_comments`, everything else is taken of by
having nodes for each branch and the usual placement.rs fixups).

## Merge Plan

This PR requires coordination between the parser repo and the main ruff
repo. I've split the ruff part, into two stacked PRs which have to be
merged together (only the second one fixes all tests), the first for the
formatter to be reviewed by @michareiser and the second for the linter
to be reviewed by @charliermarsh.

* MH: Review and merge
https://github.com/astral-sh/RustPython-Parser/pull/20
* MH: Review and merge or move later in stack
https://github.com/astral-sh/RustPython-Parser/pull/21
* MH: Review and approve
https://github.com/astral-sh/RustPython-Parser/pull/22
* MH: Review and approve formatter PR
https://github.com/astral-sh/ruff/pull/5459
* CM: Review and approve linter PR
https://github.com/astral-sh/ruff/pull/5460
* Merge linter PR in formatter PR, fix ecosystem checks (ecosystem
checks can't run on the formatter PR and won't run on the linter PR, so
we need to merge them first)
 * Merge https://github.com/astral-sh/RustPython-Parser/pull/22
 * Create tag in the parser, update linter+formatter PR
 * Merge linter+formatter PR https://github.com/astral-sh/ruff/pull/5459

---------

Co-authored-by: Micha Reiser <micha@reiser.io>

crates/ruff_python_ast/src/comparable.rs

@@ -466,6 +466,26 @@ impl<'a> From<&'a ast::ExceptHandler> for ComparableExceptHandler<'a> {
     }
 }
 
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct ComparableElifElseClause<'a> {
+    test: Option<ComparableExpr<'a>>,
+    body: Vec<ComparableStmt<'a>>,
+}
+
+impl<'a> From<&'a ast::ElifElseClause> for ComparableElifElseClause<'a> {
+    fn from(elif_else_clause: &'a ast::ElifElseClause) -> Self {
+        let ast::ElifElseClause {
+            range: _,
+            test,
+            body,
+        } = elif_else_clause;
+        Self {
+            test: test.as_ref().map(Into::into),
+            body: body.iter().map(Into::into).collect(),
+        }
+    }
+}
+
 #[derive(Debug, PartialEq, Eq, Hash)]
 pub struct ExprBoolOp<'a> {
     op: ComparableBoolOp,
@@ -999,7 +1019,7 @@ pub struct StmtWhile<'a> {
 pub struct StmtIf<'a> {
     test: ComparableExpr<'a>,
     body: Vec<ComparableStmt<'a>>,
-    orelse: Vec<ComparableStmt<'a>>,
+    elif_else_clauses: Vec<ComparableElifElseClause<'a>>,
 }
 
 #[derive(Debug, PartialEq, Eq, Hash)]
@@ -1118,7 +1138,8 @@ impl<'a> From<&'a ast::Stmt> for ComparableStmt<'a> {
                 decorator_list,
                 returns,
                 type_comment,
-                range: _range,
+                range: _,
+                type_params: _,
             }) => Self::FunctionDef(StmtFunctionDef {
                 name: name.as_str(),
                 args: args.into(),
@@ -1134,7 +1155,8 @@ impl<'a> From<&'a ast::Stmt> for ComparableStmt<'a> {
                 decorator_list,
                 returns,
                 type_comment,
-                range: _range,
+                range: _,
+                type_params: _,
             }) => Self::AsyncFunctionDef(StmtAsyncFunctionDef {
                 name: name.as_str(),
                 args: args.into(),
@@ -1149,7 +1171,8 @@ impl<'a> From<&'a ast::Stmt> for ComparableStmt<'a> {
                 keywords,
                 body,
                 decorator_list,
-                range: _range,
+                range: _,
+                type_params: _,
             }) => Self::ClassDef(StmtClassDef {
                 name: name.as_str(),
                 bases: bases.iter().map(Into::into).collect(),
@@ -1242,12 +1265,12 @@ impl<'a> From<&'a ast::Stmt> for ComparableStmt<'a> {
             ast::Stmt::If(ast::StmtIf {
                 test,
                 body,
-                orelse,
+                elif_else_clauses,
                 range: _range,
             }) => Self::If(StmtIf {
                 test: test.into(),
                 body: body.iter().map(Into::into).collect(),
-                orelse: orelse.iter().map(Into::into).collect(),
+                elif_else_clauses: elif_else_clauses.iter().map(Into::into).collect(),
             }),
             ast::Stmt::With(ast::StmtWith {
                 items,
@@ -1354,6 +1377,7 @@ impl<'a> From<&'a ast::Stmt> for ComparableStmt<'a> {
             ast::Stmt::Pass(_) => Self::Pass,
             ast::Stmt::Break(_) => Self::Break,
             ast::Stmt::Continue(_) => Self::Continue,
+            ast::Stmt::TypeAlias(_) => todo!(),
         }
     }
 }

crates/ruff_python_ast/src/helpers.rs

@@ -437,9 +437,19 @@ where
         Stmt::If(ast::StmtIf {
             test,
             body,
-            orelse,
+            elif_else_clauses,
             range: _range,
-        }) => any_over_expr(test, func) || any_over_body(body, func) || any_over_body(orelse, func),
+        }) => {
+            any_over_expr(test, func)
+                || any_over_body(body, func)
+                || elif_else_clauses.iter().any(|clause| {
+                    clause
+                        .test
+                        .as_ref()
+                        .map_or(false, |test| any_over_expr(test, func))
+                        || any_over_body(&clause.body, func)
+                })
+        }
         Stmt::With(ast::StmtWith { items, body, .. })
         | Stmt::AsyncWith(ast::StmtAsyncWith { items, body, .. }) => {
             items.iter().any(|with_item| {
@@ -529,6 +539,7 @@ where
             range: _range,
         }) => any_over_expr(value, func),
         Stmt::Pass(_) | Stmt::Break(_) | Stmt::Continue(_) => false,
+        Stmt::TypeAlias(_) => todo!(),
     }
 }
 
@@ -944,9 +955,15 @@ where
             | Stmt::AsyncFunctionDef(_)
             | Stmt::Try(_)
             | Stmt::TryStar(_) => {}
-            Stmt::If(ast::StmtIf { body, orelse, .. }) => {
+            Stmt::If(ast::StmtIf {
+                body,
+                elif_else_clauses,
+                ..
+            }) => {
                 walk_body(self, body);
-                walk_body(self, orelse);
+                for clause in elif_else_clauses {
+                    self.visit_elif_else_clause(clause);
+                }
             }
             Stmt::While(ast::StmtWhile { body, .. })
             | Stmt::With(ast::StmtWith { body, .. })
@@ -1063,25 +1080,6 @@ pub fn first_colon_range(range: TextRange, locator: &Locator) -> Option<TextRang
     range
 }
 
-/// Return the `Range` of the first `Elif` or `Else` token in an `If` statement.
-pub fn elif_else_range(stmt: &ast::StmtIf, locator: &Locator) -> Option<TextRange> {
-    let ast::StmtIf { body, orelse, .. } = stmt;
-
-    let start = body.last().expect("Expected body to be non-empty").end();
-
-    let end = match &orelse[..] {
-        [Stmt::If(ast::StmtIf { test, .. })] => test.start(),
-        [stmt, ..] => stmt.start(),
-        _ => return None,
-    };
-
-    let contents = &locator.contents()[TextRange::new(start, end)];
-    lexer::lex_starts_at(contents, Mode::Module, start)
-        .flatten()
-        .find(|(kind, _)| matches!(kind, Tok::Elif | Tok::Else))
-        .map(|(_, range)| range)
-}
-
 /// Given an offset at the end of a line (including newlines), return the offset of the
 /// continuation at the end of that line.
 fn find_continuation(offset: TextSize, locator: &Locator, indexer: &Indexer) -> Option<TextSize> {
@@ -1568,13 +1566,13 @@ mod tests {
 
     use anyhow::Result;
     use ruff_text_size::{TextLen, TextRange, TextSize};
-    use rustpython_ast::{CmpOp, Expr, Ranged, Stmt};
+    use rustpython_ast::{CmpOp, Expr, Ranged};
     use rustpython_parser::ast::Suite;
     use rustpython_parser::Parse;
 
     use crate::helpers::{
-        elif_else_range, first_colon_range, has_trailing_content, locate_cmp_ops,
-        resolve_imported_module_path, LocatedCmpOp,
+        first_colon_range, has_trailing_content, locate_cmp_ops, resolve_imported_module_path,
+        LocatedCmpOp,
     };
     use crate::source_code::Locator;
 
@@ -1667,35 +1665,6 @@ y = 2
         assert_eq!(range, TextRange::new(TextSize::from(6), TextSize::from(7)));
     }
 
-    #[test]
-    fn extract_elif_else_range() -> Result<()> {
-        let contents = "if a:
-    ...
-elif b:
-    ...
-";
-        let stmt = Stmt::parse(contents, "<filename>")?;
-        let stmt = Stmt::as_if_stmt(&stmt).unwrap();
-        let locator = Locator::new(contents);
-        let range = elif_else_range(stmt, &locator).unwrap();
-        assert_eq!(range.start(), TextSize::from(14));
-        assert_eq!(range.end(), TextSize::from(18));
-
-        let contents = "if a:
-    ...
-else:
-    ...
-";
-        let stmt = Stmt::parse(contents, "<filename>")?;
-        let stmt = Stmt::as_if_stmt(&stmt).unwrap();
-        let locator = Locator::new(contents);
-        let range = elif_else_range(stmt, &locator).unwrap();
-        assert_eq!(range.start(), TextSize::from(14));
-        assert_eq!(range.end(), TextSize::from(18));
-
-        Ok(())
-    }
-
     #[test]
     fn extract_cmp_op_location() -> Result<()> {
         let contents = "x == 1";

crates/ruff_python_ast/src/lib.rs

@@ -12,6 +12,7 @@ pub mod node;
 pub mod relocate;
 pub mod source_code;
 pub mod statement_visitor;
+pub mod stmt_if;
 pub mod str;
 pub mod token_kind;
 pub mod types;

crates/ruff_python_ast/src/node.rs

@@ -98,6 +98,7 @@ pub enum AnyNode {
     WithItem(WithItem),
     MatchCase(MatchCase),
     Decorator(Decorator),
+    ElifElseClause(ast::ElifElseClause),
 }
 
 impl AnyNode {
@@ -180,7 +181,8 @@ impl AnyNode {
             | AnyNode::Alias(_)
             | AnyNode::WithItem(_)
             | AnyNode::MatchCase(_)
-            | AnyNode::Decorator(_) => None,
+            | AnyNode::Decorator(_)
+            | AnyNode::ElifElseClause(_) => None,
         }
     }
 
@@ -263,7 +265,8 @@ impl AnyNode {
             | AnyNode::Alias(_)
             | AnyNode::WithItem(_)
             | AnyNode::MatchCase(_)
-            | AnyNode::Decorator(_) => None,
+            | AnyNode::Decorator(_)
+            | AnyNode::ElifElseClause(_) => None,
         }
     }
 
@@ -346,7 +349,8 @@ impl AnyNode {
             | AnyNode::Alias(_)
             | AnyNode::WithItem(_)
             | AnyNode::MatchCase(_)
-            | AnyNode::Decorator(_) => None,
+            | AnyNode::Decorator(_)
+            | AnyNode::ElifElseClause(_) => None,
         }
     }
 
@@ -429,7 +433,8 @@ impl AnyNode {
             | AnyNode::Alias(_)
             | AnyNode::WithItem(_)
             | AnyNode::MatchCase(_)
-            | AnyNode::Decorator(_) => None,
+            | AnyNode::Decorator(_)
+            | AnyNode::ElifElseClause(_) => None,
         }
     }
 
@@ -512,7 +517,8 @@ impl AnyNode {
             | AnyNode::Alias(_)
             | AnyNode::WithItem(_)
             | AnyNode::MatchCase(_)
-            | AnyNode::Decorator(_) => None,
+            | AnyNode::Decorator(_)
+            | AnyNode::ElifElseClause(_) => None,
         }
     }
 
@@ -595,7 +601,8 @@ impl AnyNode {
             | AnyNode::Alias(_)
             | AnyNode::WithItem(_)
             | AnyNode::MatchCase(_)
-            | AnyNode::Decorator(_) => None,
+            | AnyNode::Decorator(_)
+            | AnyNode::ElifElseClause(_) => None,
         }
     }
 
@@ -702,6 +709,7 @@ impl AnyNode {
             Self::WithItem(node) => AnyNodeRef::WithItem(node),
             Self::MatchCase(node) => AnyNodeRef::MatchCase(node),
             Self::Decorator(node) => AnyNodeRef::Decorator(node),
+            Self::ElifElseClause(node) => AnyNodeRef::ElifElseClause(node),
         }
     }
 
@@ -1159,6 +1167,34 @@ impl AstNode for ast::StmtIf {
         AnyNode::from(self)
     }
 }
+impl AstNode for ast::ElifElseClause {
+    fn cast(kind: AnyNode) -> Option<Self>
+    where
+        Self: Sized,
+    {
+        if let AnyNode::ElifElseClause(node) = kind {
+            Some(node)
+        } else {
+            None
+        }
+    }
+
+    fn cast_ref(kind: AnyNodeRef) -> Option<&Self> {
+        if let AnyNodeRef::ElifElseClause(node) = kind {
+            Some(node)
+        } else {
+            None
+        }
+    }
+
+    fn as_any_node_ref(&self) -> AnyNodeRef {
+        AnyNodeRef::from(self)
+    }
+
+    fn into_any_node(self) -> AnyNode {
+        AnyNode::from(self)
+    }
+}
 impl AstNode for ast::StmtWith {
     fn cast(kind: AnyNode) -> Option<Self>
     where
@@ -2900,6 +2936,7 @@ impl From<Stmt> for AnyNode {
             Stmt::Pass(node) => AnyNode::StmtPass(node),
             Stmt::Break(node) => AnyNode::StmtBreak(node),
             Stmt::Continue(node) => AnyNode::StmtContinue(node),
+            Stmt::TypeAlias(_) => todo!(),
         }
     }
 }
@@ -3076,6 +3113,12 @@ impl From<ast::StmtIf> for AnyNode {
     }
 }
 
+impl From<ast::ElifElseClause> for AnyNode {
+    fn from(node: ast::ElifElseClause) -> Self {
+        AnyNode::ElifElseClause(node)
+    }
+}
+
 impl From<ast::StmtWith> for AnyNode {
     fn from(node: ast::StmtWith) -> Self {
         AnyNode::StmtWith(node)
@@ -3514,6 +3557,7 @@ impl Ranged for AnyNode {
             AnyNode::WithItem(node) => node.range(),
             AnyNode::MatchCase(node) => node.range(),
             AnyNode::Decorator(node) => node.range(),
+            AnyNode::ElifElseClause(node) => node.range(),
         }
     }
 }
@@ -3597,6 +3641,7 @@ pub enum AnyNodeRef<'a> {
     WithItem(&'a WithItem),
     MatchCase(&'a MatchCase),
     Decorator(&'a Decorator),
+    ElifElseClause(&'a ast::ElifElseClause),
 }
 
 impl AnyNodeRef<'_> {
@@ -3679,6 +3724,7 @@ impl AnyNodeRef<'_> {
             AnyNodeRef::WithItem(node) => NonNull::from(*node).cast(),
             AnyNodeRef::MatchCase(node) => NonNull::from(*node).cast(),
             AnyNodeRef::Decorator(node) => NonNull::from(*node).cast(),
+            AnyNodeRef::ElifElseClause(node) => NonNull::from(*node).cast(),
         }
     }
 
@@ -3767,6 +3813,7 @@ impl AnyNodeRef<'_> {
             AnyNodeRef::WithItem(_) => NodeKind::WithItem,
             AnyNodeRef::MatchCase(_) => NodeKind::MatchCase,
             AnyNodeRef::Decorator(_) => NodeKind::Decorator,
+            AnyNodeRef::ElifElseClause(_) => NodeKind::ElifElseClause,
         }
     }
 
@@ -3849,7 +3896,8 @@ impl AnyNodeRef<'_> {
             | AnyNodeRef::Alias(_)
             | AnyNodeRef::WithItem(_)
             | AnyNodeRef::MatchCase(_)
-            | AnyNodeRef::Decorator(_) => false,
+            | AnyNodeRef::Decorator(_)
+            | AnyNodeRef::ElifElseClause(_) => false,
         }
     }
 
@@ -3932,7 +3980,8 @@ impl AnyNodeRef<'_> {
             | AnyNodeRef::Alias(_)
             | AnyNodeRef::WithItem(_)
             | AnyNodeRef::MatchCase(_)
-            | AnyNodeRef::Decorator(_) => false,
+            | AnyNodeRef::Decorator(_)
+            | AnyNodeRef::ElifElseClause(_) => false,
         }
     }
 
@@ -4015,7 +4064,8 @@ impl AnyNodeRef<'_> {
             | AnyNodeRef::Alias(_)
             | AnyNodeRef::WithItem(_)
             | AnyNodeRef::MatchCase(_)
-            | AnyNodeRef::Decorator(_) => false,
+            | AnyNodeRef::Decorator(_)
+            | AnyNodeRef::ElifElseClause(_) => false,
         }
     }
 
@@ -4098,7 +4148,8 @@ impl AnyNodeRef<'_> {
             | AnyNodeRef::Alias(_)
             | AnyNodeRef::WithItem(_)
             | AnyNodeRef::MatchCase(_)
-            | AnyNodeRef::Decorator(_) => false,
+            | AnyNodeRef::Decorator(_)
+            | AnyNodeRef::ElifElseClause(_) => false,
         }
     }
 
@@ -4181,7 +4232,8 @@ impl AnyNodeRef<'_> {
             | AnyNodeRef::Alias(_)
             | AnyNodeRef::WithItem(_)
             | AnyNodeRef::MatchCase(_)
-            | AnyNodeRef::Decorator(_) => false,
+            | AnyNodeRef::Decorator(_)
+            | AnyNodeRef::ElifElseClause(_) => false,
         }
     }
 
@@ -4264,7 +4316,8 @@ impl AnyNodeRef<'_> {
             | AnyNodeRef::Alias(_)
             | AnyNodeRef::WithItem(_)
             | AnyNodeRef::MatchCase(_)
-            | AnyNodeRef::Decorator(_) => false,
+            | AnyNodeRef::Decorator(_)
+            | AnyNodeRef::ElifElseClause(_) => false,
         }
     }
 
@@ -4383,6 +4436,12 @@ impl<'a> From<&'a ast::StmtIf> for AnyNodeRef<'a> {
     }
 }
 
+impl<'a> From<&'a ast::ElifElseClause> for AnyNodeRef<'a> {
+    fn from(node: &'a ast::ElifElseClause) -> Self {
+        AnyNodeRef::ElifElseClause(node)
+    }
+}
+
 impl<'a> From<&'a ast::StmtWith> for AnyNodeRef<'a> {
     fn from(node: &'a ast::StmtWith) -> Self {
         AnyNodeRef::StmtWith(node)
@@ -4731,6 +4790,7 @@ impl<'a> From<&'a Stmt> for AnyNodeRef<'a> {
             Stmt::Pass(node) => AnyNodeRef::StmtPass(node),
             Stmt::Break(node) => AnyNodeRef::StmtBreak(node),
             Stmt::Continue(node) => AnyNodeRef::StmtContinue(node),
+            Stmt::TypeAlias(_) => todo!(),
         }
     }
 }
@@ -4934,6 +4994,7 @@ impl Ranged for AnyNodeRef<'_> {
             AnyNodeRef::WithItem(node) => node.range(),
             AnyNodeRef::MatchCase(node) => node.range(),
             AnyNodeRef::Decorator(node) => node.range(),
+            AnyNodeRef::ElifElseClause(node) => node.range(),
         }
     }
 }
@@ -5017,4 +5078,5 @@ pub enum NodeKind {
     WithItem,
     MatchCase,
     Decorator,
+    ElifElseClause,
 }

crates/ruff_python_ast/src/source_code/comment_ranges.rs

@@ -1,3 +1,4 @@
+use itertools::Itertools;
 use std::fmt::{Debug, Formatter};
 use std::ops::Deref;
 
@@ -25,6 +26,21 @@ impl CommentRanges {
             })
             .is_ok()
     }
+
+    /// Returns the comments who are within the range
+    pub fn comments_in_range(&self, range: TextRange) -> &[TextRange] {
+        let start = self
+            .raw
+            .partition_point(|comment| comment.start() < range.start());
+        // We expect there are few comments, so switching to find should be faster
+        match self.raw[start..]
+            .iter()
+            .find_position(|comment| comment.end() > range.end())
+        {
+            Some((in_range, _element)) => &self.raw[start..start + in_range],
+            None => &self.raw[start..],
+        }
+    }
 }
 
 impl Deref for CommentRanges {

crates/ruff_python_ast/src/source_code/generator.rs

@@ -271,7 +271,8 @@ impl<'a> Generator<'a> {
                 keywords,
                 body,
                 decorator_list,
-                range: _range,
+                range: _,
+                type_params: _,
             }) => {
                 self.newlines(if self.indent_depth == 0 { 2 } else { 1 });
                 for decorator in decorator_list {
@@ -457,7 +458,7 @@ impl<'a> Generator<'a> {
             Stmt::If(ast::StmtIf {
                 test,
                 body,
-                orelse,
+                elif_else_clauses,
                 range: _range,
             }) => {
                 statement!({
@@ -467,33 +468,19 @@ impl<'a> Generator<'a> {
                 });
                 self.body(body);
 
-                let mut orelse_: &[Stmt] = orelse;
-                loop {
-                    if orelse_.len() == 1 && matches!(orelse_[0], Stmt::If(_)) {
-                        if let Stmt::If(ast::StmtIf {
-                            body,
-                            test,
-                            orelse,
-                            range: _range,
-                        }) = &orelse_[0]
-                        {
-                            statement!({
-                                self.p("elif ");
-                                self.unparse_expr(test, precedence::IF);
-                                self.p(":");
-                            });
-                            self.body(body);
-                            orelse_ = orelse;
-                        }
+                for clause in elif_else_clauses {
+                    if let Some(test) = &clause.test {
+                        statement!({
+                            self.p("elif ");
+                            self.unparse_expr(test, precedence::IF);
+                            self.p(":");
+                        });
                     } else {
-                        if !orelse_.is_empty() {
-                            statement!({
-                                self.p("else:");
-                            });
-                            self.body(orelse_);
-                        }
-                        break;
+                        statement!({
+                            self.p("else:");
+                        });
                     }
+                    self.body(&clause.body);
                 }
             }
             Stmt::With(ast::StmtWith { items, body, .. }) => {
@@ -715,6 +702,7 @@ impl<'a> Generator<'a> {
                     self.p("continue");
                 });
             }
+            Stmt::TypeAlias(_) => todo!(),
         }
     }
 

crates/ruff_python_ast/src/source_code/indexer.rs

@@ -97,6 +97,18 @@ impl Indexer {
         &self.comment_ranges
     }
 
+    /// Returns the comments in the given range as source code slices
+    pub fn comments_in_range<'a>(
+        &'a self,
+        range: TextRange,
+        locator: &'a Locator,
+    ) -> impl Iterator<Item = &'a str> {
+        self.comment_ranges
+            .comments_in_range(range)
+            .iter()
+            .map(move |comment_range| locator.slice(*comment_range))
+    }
+
     /// Returns the line start positions of continuations (backslash).
     pub fn continuation_line_starts(&self) -> &[TextSize] {
         &self.continuation_lines

crates/ruff_python_ast/src/statement_visitor.rs

@@ -1,5 +1,6 @@
 //! Specialized AST visitor trait and walk functions that only visit statements.
 
+use rustpython_ast::ElifElseClause;
 use rustpython_parser::ast::{self, ExceptHandler, MatchCase, Stmt};
 
 /// A trait for AST visitors that only need to visit statements.
@@ -13,6 +14,9 @@ pub trait StatementVisitor<'a> {
     fn visit_except_handler(&mut self, except_handler: &'a ExceptHandler) {
         walk_except_handler(self, except_handler);
     }
+    fn visit_elif_else_clause(&mut self, elif_else_clause: &'a ElifElseClause) {
+        walk_elif_else_clause(self, elif_else_clause);
+    }
     fn visit_match_case(&mut self, match_case: &'a MatchCase) {
         walk_match_case(self, match_case);
     }
@@ -47,9 +51,15 @@ pub fn walk_stmt<'a, V: StatementVisitor<'a> + ?Sized>(visitor: &mut V, stmt: &'
             visitor.visit_body(body);
             visitor.visit_body(orelse);
         }
-        Stmt::If(ast::StmtIf { body, orelse, .. }) => {
+        Stmt::If(ast::StmtIf {
+            body,
+            elif_else_clauses,
+            ..
+        }) => {
             visitor.visit_body(body);
-            visitor.visit_body(orelse);
+            for clause in elif_else_clauses {
+                visitor.visit_elif_else_clause(clause);
+            }
         }
         Stmt::With(ast::StmtWith { body, .. }) => {
             visitor.visit_body(body);
@@ -105,6 +115,13 @@ pub fn walk_except_handler<'a, V: StatementVisitor<'a> + ?Sized>(
     }
 }
 
+pub fn walk_elif_else_clause<'a, V: StatementVisitor<'a> + ?Sized>(
+    visitor: &mut V,
+    elif_else_clause: &'a ElifElseClause,
+) {
+    visitor.visit_body(&elif_else_clause.body);
+}
+
 pub fn walk_match_case<'a, V: StatementVisitor<'a> + ?Sized>(
     visitor: &mut V,
     match_case: &'a MatchCase,

crates/ruff_python_ast/src/stmt_if.rs

@@ -0,0 +1,87 @@
+use crate::source_code::Locator;
+use ruff_text_size::TextRange;
+use rustpython_ast::{ElifElseClause, Expr, Ranged, Stmt, StmtIf};
+use rustpython_parser::{lexer, Mode, Tok};
+use std::iter;
+
+/// Return the `Range` of the first `Elif` or `Else` token in an `If` statement.
+pub fn elif_else_range(clause: &ElifElseClause, locator: &Locator) -> Option<TextRange> {
+    let contents = &locator.contents()[clause.range];
+    let token = lexer::lex_starts_at(contents, Mode::Module, clause.range.start())
+        .flatten()
+        .next()?;
+    if matches!(token.0, Tok::Elif | Tok::Else) {
+        Some(token.1)
+    } else {
+        None
+    }
+}
+
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub enum BranchKind {
+    If,
+    Elif,
+}
+
+pub struct IfElifBranch<'a> {
+    pub kind: BranchKind,
+    pub test: &'a Expr,
+    pub body: &'a [Stmt],
+    pub range: TextRange,
+}
+
+pub fn if_elif_branches(stmt_if: &StmtIf) -> impl Iterator<Item = IfElifBranch> {
+    iter::once(IfElifBranch {
+        kind: BranchKind::If,
+        test: stmt_if.test.as_ref(),
+        body: stmt_if.body.as_slice(),
+        range: TextRange::new(stmt_if.start(), stmt_if.body.last().unwrap().end()),
+    })
+    .chain(stmt_if.elif_else_clauses.iter().filter_map(|clause| {
+        Some(IfElifBranch {
+            kind: BranchKind::Elif,
+            test: clause.test.as_ref()?,
+            body: clause.body.as_slice(),
+            range: clause.range,
+        })
+    }))
+}
+
+#[cfg(test)]
+mod test {
+    use crate::source_code::Locator;
+    use crate::stmt_if::elif_else_range;
+    use anyhow::Result;
+    use ruff_text_size::TextSize;
+    use rustpython_ast::Stmt;
+    use rustpython_parser::Parse;
+
+    #[test]
+    fn extract_elif_else_range() -> Result<()> {
+        let contents = "if a:
+    ...
+elif b:
+    ...
+";
+        let stmt = Stmt::parse(contents, "<filename>")?;
+        let stmt = Stmt::as_if_stmt(&stmt).unwrap();
+        let locator = Locator::new(contents);
+        let range = elif_else_range(&stmt.elif_else_clauses[0], &locator).unwrap();
+        assert_eq!(range.start(), TextSize::from(14));
+        assert_eq!(range.end(), TextSize::from(18));
+
+        let contents = "if a:
+    ...
+else:
+    ...
+";
+        let stmt = Stmt::parse(contents, "<filename>")?;
+        let stmt = Stmt::as_if_stmt(&stmt).unwrap();
+        let locator = Locator::new(contents);
+        let range = elif_else_range(&stmt.elif_else_clauses[0], &locator).unwrap();
+        assert_eq!(range.start(), TextSize::from(14));
+        assert_eq!(range.end(), TextSize::from(18));
+
+        Ok(())
+    }
+}

crates/ruff_python_ast/src/token_kind.rs

@@ -431,6 +431,8 @@ impl TokenKind {
             Tok::StartModule => TokenKind::StartModule,
             Tok::StartInteractive => TokenKind::StartInteractive,
             Tok::StartExpression => TokenKind::StartExpression,
+            Tok::MagicCommand { .. } => todo!(),
+            Tok::Type => todo!(),
         }
     }
 }

crates/ruff_python_ast/src/visitor.rs

@@ -2,6 +2,7 @@
 
 pub mod preorder;
 
+use rustpython_ast::ElifElseClause;
 use rustpython_parser::ast::{
     self, Alias, Arg, Arguments, BoolOp, CmpOp, Comprehension, Decorator, ExceptHandler, Expr,
     ExprContext, Keyword, MatchCase, Operator, Pattern, Stmt, UnaryOp, WithItem,
@@ -75,6 +76,9 @@ pub trait Visitor<'a> {
     fn visit_body(&mut self, body: &'a [Stmt]) {
         walk_body(self, body);
     }
+    fn visit_elif_else_clause(&mut self, elif_else_clause: &'a ElifElseClause) {
+        walk_elif_else_clause(self, elif_else_clause);
+    }
 }
 
 pub fn walk_body<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, body: &'a [Stmt]) {
@@ -83,6 +87,16 @@ pub fn walk_body<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, body: &'a [Stmt])
     }
 }
 
+pub fn walk_elif_else_clause<'a, V: Visitor<'a> + ?Sized>(
+    visitor: &mut V,
+    elif_else_clause: &'a ElifElseClause,
+) {
+    if let Some(test) = &elif_else_clause.test {
+        visitor.visit_expr(test);
+    }
+    visitor.visit_body(&elif_else_clause.body);
+}
+
 pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) {
     match stmt {
         Stmt::FunctionDef(ast::StmtFunctionDef {
@@ -216,12 +230,17 @@ pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) {
         Stmt::If(ast::StmtIf {
             test,
             body,
-            orelse,
+            elif_else_clauses,
             range: _range,
         }) => {
             visitor.visit_expr(test);
             visitor.visit_body(body);
-            visitor.visit_body(orelse);
+            for clause in elif_else_clauses {
+                if let Some(test) = &clause.test {
+                    visitor.visit_expr(test);
+                }
+                walk_elif_else_clause(visitor, clause);
+            }
         }
         Stmt::With(ast::StmtWith { items, body, .. }) => {
             for with_item in items {
@@ -315,6 +334,7 @@ pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) {
             range: _range,
         }) => visitor.visit_expr(value),
         Stmt::Pass(_) | Stmt::Break(_) | Stmt::Continue(_) => {}
+        Stmt::TypeAlias(_) => todo!(),
     }
 }
 

crates/ruff_python_ast/src/visitor/preorder.rs

@@ -1,4 +1,4 @@
-use rustpython_ast::{ArgWithDefault, Mod, TypeIgnore};
+use rustpython_ast::{ArgWithDefault, ElifElseClause, Mod, TypeIgnore};
 use rustpython_parser::ast::{
     self, Alias, Arg, Arguments, BoolOp, CmpOp, Comprehension, Constant, Decorator, ExceptHandler,
     Expr, Keyword, MatchCase, Operator, Pattern, Stmt, UnaryOp, WithItem,
@@ -95,6 +95,10 @@ pub trait PreorderVisitor<'a> {
     fn visit_type_ignore(&mut self, type_ignore: &'a TypeIgnore) {
         walk_type_ignore(self, type_ignore);
     }
+
+    fn visit_elif_else_clause(&mut self, elif_else_clause: &'a ElifElseClause) {
+        walk_elif_else_clause(self, elif_else_clause);
+    }
 }
 
 pub fn walk_module<'a, V>(visitor: &mut V, module: &'a Mod)
@@ -286,12 +290,14 @@ where
         Stmt::If(ast::StmtIf {
             test,
             body,
-            orelse,
+            elif_else_clauses,
             range: _range,
         }) => {
             visitor.visit_expr(test);
             visitor.visit_body(body);
-            visitor.visit_body(orelse);
+            for clause in elif_else_clauses {
+                visitor.visit_elif_else_clause(clause);
+            }
         }
 
         Stmt::With(ast::StmtWith {
@@ -394,6 +400,7 @@ where
         | Stmt::Continue(_)
         | Stmt::Global(_)
         | Stmt::Nonlocal(_) => {}
+        Stmt::TypeAlias(_) => todo!(),
     }
 }
 
@@ -703,6 +710,16 @@ where
     }
 }
 
+pub fn walk_elif_else_clause<'a, V>(visitor: &mut V, elif_else_clause: &'a ElifElseClause)
+where
+    V: PreorderVisitor<'a> + ?Sized,
+{
+    if let Some(test) = &elif_else_clause.test {
+        visitor.visit_expr(test);
+    }
+    visitor.visit_body(&elif_else_clause.body);
+}
+
 pub fn walk_except_handler<'a, V>(visitor: &mut V, except_handler: &'a ExceptHandler)
 where
     V: PreorderVisitor<'a> + ?Sized,