Range formatting API (#9635)

crates/ruff_python_formatter/src/statement/clause.rs

@@ -357,8 +357,20 @@ impl<'ast> Format<PyFormatContext<'ast>> for FormatClauseHeader<'_, 'ast> {
         if SuppressionKind::has_skip_comment(self.trailing_colon_comment, f.context().source()) {
             write_suppressed_clause_header(self.header, f)?;
         } else {
-            f.write_fmt(Arguments::from(&self.formatter))?;
-            token(":").fmt(f)?;
+            // Write a source map entry for the colon for range formatting to support formatting the clause header without
+            // the clause body. Avoid computing `self.header.range()` otherwise because it's somewhat involved.
+            let clause_end = if f.options().source_map_generation().is_enabled() {
+                Some(source_position(
+                    self.header.range(f.context().source())?.end(),
+                ))
+            } else {
+                None
+            };
+
+            write!(
+                f,
+                [Arguments::from(&self.formatter), token(":"), clause_end]
+            )?;
         }
 
         trailing_comments(self.trailing_colon_comment).fmt(f)
@@ -458,7 +470,12 @@ fn find_keyword(
 fn colon_range(after_keyword_or_condition: TextSize, source: &str) -> FormatResult<TextRange> {
     let mut tokenizer = SimpleTokenizer::starts_at(after_keyword_or_condition, source)
         .skip_trivia()
-        .skip_while(|token| token.kind() == SimpleTokenKind::RParen);
+        .skip_while(|token| {
+            matches!(
+                token.kind(),
+                SimpleTokenKind::RParen | SimpleTokenKind::Comma
+            )
+        });
 
     match tokenizer.next() {
         Some(SimpleToken {

crates/ruff_python_formatter/src/statement/stmt_if.rs

@@ -1,6 +1,7 @@
 use ruff_formatter::{format_args, write};
 use ruff_python_ast::AnyNodeRef;
 use ruff_python_ast::{ElifElseClause, StmtIf};
+use ruff_text_size::Ranged;
 
 use crate::comments::SourceComment;
 use crate::expression::maybe_parenthesize_expression;
@@ -81,7 +82,12 @@ pub(crate) fn format_elif_else_clause(
             clause_header(
                 ClauseHeader::ElifElse(item),
                 trailing_colon_comment,
-                &format_with(|f| {
+                &format_with(|f: &mut PyFormatter| {
+                    f.options()
+                        .source_map_generation()
+                        .is_enabled()
+                        .then_some(source_position(item.start()))
+                        .fmt(f)?;
                     if let Some(test) = test {
                         write!(
                             f,
@@ -98,6 +104,10 @@ pub(crate) fn format_elif_else_clause(
             )
             .with_leading_comments(leading_comments, last_node),
             clause_body(body, trailing_colon_comment),
+            f.options()
+                .source_map_generation()
+                .is_enabled()
+                .then_some(source_position(item.end()))
         ]
     )
 }

crates/ruff_python_formatter/src/statement/suite.rs

@@ -2,8 +2,8 @@ use ruff_formatter::{
     write, FormatContext, FormatOwnedWithRule, FormatRefWithRule, FormatRuleWithOptions,
 };
 use ruff_python_ast::helpers::is_compound_statement;
-use ruff_python_ast::AnyNodeRef;
 use ruff_python_ast::{self as ast, Expr, PySourceType, Stmt, Suite};
+use ruff_python_ast::{AnyNodeRef, StmtExpr};
 use ruff_python_trivia::{lines_after, lines_after_ignoring_end_of_line_trivia, lines_before};
 use ruff_text_size::{Ranged, TextRange};
 
@@ -740,6 +740,14 @@ impl<'a> DocstringStmt<'a> {
             _ => None,
         }
     }
+
+    pub(crate) fn is_docstring_statement(stmt: &StmtExpr) -> bool {
+        if let Expr::StringLiteral(ast::ExprStringLiteral { value, .. }) = stmt.value.as_ref() {
+            !value.is_implicit_concatenated()
+        } else {
+            false
+        }
+    }
 }
 
 impl Format<PyFormatContext<'_>> for DocstringStmt<'_> {