Range formatting API (#9635)

crates/ruff_formatter/src/builders.rs

@@ -328,6 +328,12 @@ pub struct SourcePosition(TextSize);
 
 impl<Context> Format<Context> for SourcePosition {
     fn fmt(&self, f: &mut Formatter<Context>) -> FormatResult<()> {
+        if let Some(FormatElement::SourcePosition(last_position)) = f.buffer.elements().last() {
+            if *last_position == self.0 {
+                return Ok(());
+            }
+        }
+
         f.write_element(FormatElement::SourcePosition(self.0));
 
         Ok(())
@@ -353,8 +359,8 @@ where
     Context: FormatContext,
 {
     fn fmt(&self, f: &mut Formatter<Context>) -> FormatResult<()> {
-        if let Some(source_position) = self.position {
-            f.write_element(FormatElement::SourcePosition(source_position));
+        if let Some(position) = self.position {
+            source_position(position).fmt(f)?;
         }
 
         f.write_element(FormatElement::Text {

crates/ruff_formatter/src/lib.rs

@@ -53,7 +53,7 @@ pub use crate::diagnostics::{ActualStart, FormatError, InvalidDocumentError, Pri
 pub use format_element::{normalize_newlines, FormatElement, LINE_TERMINATORS};
 pub use group_id::GroupId;
 use ruff_macros::CacheKey;
-use ruff_text_size::{TextRange, TextSize};
+use ruff_text_size::{TextLen, TextRange, TextSize};
 
 #[derive(Debug, Eq, PartialEq, Clone, Copy, Hash, CacheKey)]
 #[cfg_attr(
@@ -432,6 +432,90 @@ impl Printed {
     pub fn take_verbatim_ranges(&mut self) -> Vec<TextRange> {
         std::mem::take(&mut self.verbatim_ranges)
     }
+
+    /// Slices the formatted code to the sub-slices that covers the passed `source_range`.
+    ///
+    /// The implementation uses the source map generated during formatting to find the closest range
+    /// in the formatted document that covers `source_range` or more. The returned slice
+    /// matches the `source_range` exactly (except indent, see below) if the formatter emits [`FormatElement::SourcePosition`] for
+    /// the range's offsets.
+    ///
+    /// Returns the entire document if the source map is empty.
+    #[must_use]
+    pub fn slice_range(self, source_range: TextRange) -> PrintedRange {
+        let mut start_marker: Option<SourceMarker> = None;
+        let mut end_marker: Option<SourceMarker> = None;
+
+        // Note: The printer can generate multiple source map entries for the same source position.
+        // For example if you have:
+        // * `source_position(276)`
+        // * `token("def")`
+        // * `token("foo")`
+        // * `source_position(284)`
+        // The printer uses the source position 276 for both the tokens `def` and `foo` because that's the only position it knows of.
+        //
+        // Warning: Source markers are often emitted sorted by their source position but it's not guaranteed.
+        // They are only guaranteed to be sorted in increasing order by their destination position.
+        for marker in self.sourcemap {
+            // Take the closest start marker, but skip over start_markers that have the same start.
+            if marker.source <= source_range.start()
+                && !start_marker.is_some_and(|existing| existing.source >= marker.source)
+            {
+                start_marker = Some(marker);
+            }
+
+            if marker.source >= source_range.end()
+                && !end_marker.is_some_and(|existing| existing.source <= marker.source)
+            {
+                end_marker = Some(marker);
+            }
+        }
+
+        let start = start_marker.map(|marker| marker.dest).unwrap_or_default();
+        let end = end_marker.map_or_else(|| self.code.text_len(), |marker| marker.dest);
+        let code_range = TextRange::new(start, end);
+
+        PrintedRange {
+            code: self.code[code_range].to_string(),
+            source_range,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
+pub struct PrintedRange {
+    code: String,
+    source_range: TextRange,
+}
+
+impl PrintedRange {
+    pub fn new(code: String, source_range: TextRange) -> Self {
+        Self { code, source_range }
+    }
+
+    pub fn empty() -> Self {
+        Self {
+            code: String::new(),
+            source_range: TextRange::default(),
+        }
+    }
+
+    /// The formatted code.
+    pub fn as_code(&self) -> &str {
+        &self.code
+    }
+
+    /// The range the formatted code corresponds to in the source document.
+    pub fn source_range(&self) -> TextRange {
+        self.source_range
+    }
+
+    #[must_use]
+    pub fn with_code(self, code: String) -> Self {
+        Self { code, ..self }
+    }
 }
 
 /// Public return type of the formatter

crates/ruff_formatter/src/printer/mod.rs

@@ -60,7 +60,10 @@ impl<'a> Printer<'a> {
         document: &'a Document,
         indent: u16,
     ) -> PrintResult<Printed> {
-        let mut stack = PrintCallStack::new(PrintElementArgs::new(Indention::Level(indent)));
+        let indentation = Indention::Level(indent);
+        self.state.pending_indent = indentation;
+
+        let mut stack = PrintCallStack::new(PrintElementArgs::new(indentation));
         let mut queue: PrintQueue<'a> = PrintQueue::new(document.as_ref());
 
         loop {
@@ -143,7 +146,13 @@ impl<'a> Printer<'a> {
 
             FormatElement::SourcePosition(position) => {
                 self.state.source_position = *position;
-                self.push_marker();
+                // The printer defers printing indents until the next text
+                // is printed. Pushing the marker now would mean that the
+                // mapped range includes the indent range, which we don't want.
+                // Only add a marker if we're not in an indented context, e.g. at the end of the file.
+                if self.state.pending_indent.is_empty() {
+                    self.push_marker();
+                }
             }
 
             FormatElement::LineSuffixBoundary => {
@@ -511,11 +520,12 @@ impl<'a> Printer<'a> {
             dest: self.state.buffer.text_len(),
         };
 
-        if let Some(last) = self.state.source_markers.last() {
-            if last != &marker {
-                self.state.source_markers.push(marker);
-            }
-        } else {
+        if self
+            .state
+            .source_markers
+            .last()
+            .map_or(true, |last| last != &marker)
+        {
             self.state.source_markers.push(marker);
         }
     }