Format if statements (#4961)

crates/ruff_python_formatter/src/statement/stmt_if.rs

@@ -1,12 +1,130 @@
-use crate::{not_yet_implemented, FormatNodeRule, PyFormatter};
-use ruff_formatter::{write, Buffer, FormatResult};
-use rustpython_parser::ast::StmtIf;
+use crate::comments::{leading_alternate_branch_comments, trailing_comments, SourceComment};
+use crate::expression::parentheses::Parenthesize;
+use crate::prelude::*;
+use crate::FormatNodeRule;
+use ruff_formatter::{write, FormatError};
+use rustpython_parser::ast::{Ranged, Stmt, StmtIf, Suite};
 
 #[derive(Default)]
 pub struct FormatStmtIf;
 
 impl FormatNodeRule<StmtIf> for FormatStmtIf {
     fn fmt_fields(&self, item: &StmtIf, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(f, [not_yet_implemented(item)])
+        let comments = f.context().comments().clone();
+
+        let mut current = IfOrElIf::If(item);
+        let mut else_comments: &[SourceComment];
+        let mut last_node_of_previous_body = None;
+
+        loop {
+            let current_statement = current.statement();
+            let StmtIf {
+                test, body, orelse, ..
+            } = current_statement;
+
+            let first_statement = body.first().ok_or(FormatError::SyntaxError)?;
+            let trailing = comments.dangling_comments(current_statement.into());
+
+            let trailing_if_comments_end = trailing
+                .partition_point(|comment| comment.slice().start() < first_statement.start());
+
+            let (if_trailing_comments, trailing_alternate_comments) =
+                trailing.split_at(trailing_if_comments_end);
+
+            if current.is_elif() {
+                let elif_leading = comments.leading_comments(current_statement.into());
+                // Manually format the leading comments because the formatting bypasses `NodeRule::fmt`
+                write!(
+                    f,
+                    [
+                        leading_alternate_branch_comments(elif_leading, last_node_of_previous_body),
+                        source_position(current_statement.start())
+                    ]
+                )?;
+            }
+
+            write!(
+                f,
+                [
+                    text(current.keyword()),
+                    space(),
+                    test.format().with_options(Parenthesize::IfBreaks),
+                    text(":"),
+                    trailing_comments(if_trailing_comments),
+                    block_indent(&body.format())
+                ]
+            )?;
+
+            // RustPython models `elif` by setting the body to a single `if` statement. The `orelse`
+            // of the most inner `if` statement then becomes the `else` of the whole `if` chain.
+            // That's why it's necessary to take the comments here from the most inner `elif`.
+            else_comments = trailing_alternate_comments;
+            last_node_of_previous_body = body.last();
+
+            if let Some(elif) = else_if(orelse) {
+                current = elif;
+            } else {
+                break;
+            }
+        }
+
+        let orelse = &current.statement().orelse;
+
+        if !orelse.is_empty() {
+            // Leading comments are always own line comments
+            let leading_else_comments_end =
+                else_comments.partition_point(|comment| comment.position().is_own_line());
+            let (else_leading, else_trailing) = else_comments.split_at(leading_else_comments_end);
+
+            write!(
+                f,
+                [
+                    leading_alternate_branch_comments(else_leading, last_node_of_previous_body),
+                    text("else:"),
+                    trailing_comments(else_trailing),
+                    block_indent(&orelse.format())
+                ]
+            )?;
+        }
+
+        Ok(())
+    }
+
+    fn fmt_dangling_comments(&self, _node: &StmtIf, _f: &mut PyFormatter) -> FormatResult<()> {
+        // Handled by `fmt_fields`
+        Ok(())
+    }
+}
+
+fn else_if(or_else: &Suite) -> Option<IfOrElIf> {
+    if let [Stmt::If(if_stmt)] = or_else.as_slice() {
+        Some(IfOrElIf::ElIf(if_stmt))
+    } else {
+        None
+    }
+}
+
+enum IfOrElIf<'a> {
+    If(&'a StmtIf),
+    ElIf(&'a StmtIf),
+}
+
+impl<'a> IfOrElIf<'a> {
+    const fn statement(&self) -> &'a StmtIf {
+        match self {
+            IfOrElIf::If(statement) => statement,
+            IfOrElIf::ElIf(statement) => statement,
+        }
+    }
+
+    const fn keyword(&self) -> &'static str {
+        match self {
+            IfOrElIf::If(_) => "if",
+            IfOrElIf::ElIf(_) => "elif",
+        }
+    }
+
+    const fn is_elif(&self) -> bool {
+        matches!(self, IfOrElIf::ElIf(_))
     }
 }