Range formatting API (#9635)

crates/ruff_python_formatter/src/comments/format.rs

@@ -432,7 +432,17 @@ impl Format<PyFormatContext<'_>> for FormatNormalizedComment<'_> {
     fn fmt(&self, f: &mut Formatter<PyFormatContext>) -> FormatResult<()> {
         match self.comment {
             Cow::Borrowed(borrowed) => {
-                source_text_slice(TextRange::at(self.range.start(), borrowed.text_len())).fmt(f)
+                source_text_slice(TextRange::at(self.range.start(), borrowed.text_len())).fmt(f)?;
+
+                // Write the end position if the borrowed comment is shorter than the original comment
+                // So that we still can map back the end of a comment to the formatted code.
+                if f.options().source_map_generation().is_enabled()
+                    && self.range.len() != borrowed.text_len()
+                {
+                    source_position(self.range.end()).fmt(f)?;
+                }
+
+                Ok(())
             }
 
             Cow::Owned(ref owned) => {

crates/ruff_python_formatter/src/comments/mod.rs

@@ -98,7 +98,6 @@ pub(crate) use format::{
 use ruff_formatter::{SourceCode, SourceCodeSlice};
 use ruff_python_ast::visitor::preorder::{PreorderVisitor, TraversalSignal};
 use ruff_python_ast::AnyNodeRef;
-use ruff_python_ast::Mod;
 use ruff_python_trivia::{CommentRanges, PythonWhitespace};
 use ruff_source_file::Locator;
 use ruff_text_size::{Ranged, TextRange};
@@ -347,20 +346,28 @@ impl<'a> Comments<'a> {
 
     /// Extracts the comments from the AST.
     pub(crate) fn from_ast(
-        root: &'a Mod,
+        root: impl Into<AnyNodeRef<'a>>,
         source_code: SourceCode<'a>,
         comment_ranges: &'a CommentRanges,
     ) -> Self {
-        let map = if comment_ranges.is_empty() {
-            CommentsMap::new()
-        } else {
-            let mut builder =
-                CommentsMapBuilder::new(Locator::new(source_code.as_str()), comment_ranges);
-            CommentsVisitor::new(source_code, comment_ranges, &mut builder).visit(root);
-            builder.finish()
-        };
+        fn collect_comments<'a>(
+            root: AnyNodeRef<'a>,
+            source_code: SourceCode<'a>,
+            comment_ranges: &'a CommentRanges,
+        ) -> Comments<'a> {
+            let map = if comment_ranges.is_empty() {
+                CommentsMap::new()
+            } else {
+                let mut builder =
+                    CommentsMapBuilder::new(Locator::new(source_code.as_str()), comment_ranges);
+                CommentsVisitor::new(source_code, comment_ranges, &mut builder).visit(root);
+                builder.finish()
+            };
 
-        Self::new(map, comment_ranges)
+            Comments::new(map, comment_ranges)
+        }
+
+        collect_comments(root.into(), source_code, comment_ranges)
     }
 
     /// Returns `true` if the given `node` has any comments.

crates/ruff_python_formatter/src/comments/visitor.rs

@@ -24,7 +24,7 @@ pub(crate) fn collect_comments<'a>(
     comment_ranges: &'a CommentRanges,
 ) -> Vec<DecoratedComment<'a>> {
     let mut collector = CommentsVecBuilder::default();
-    CommentsVisitor::new(source_code, comment_ranges, &mut collector).visit(root);
+    CommentsVisitor::new(source_code, comment_ranges, &mut collector).visit(AnyNodeRef::from(root));
     collector.comments
 }
 
@@ -52,8 +52,12 @@ impl<'a, 'builder> CommentsVisitor<'a, 'builder> {
         }
     }
 
-    pub(super) fn visit(mut self, root: &'a Mod) {
-        self.visit_mod(root);
+    pub(super) fn visit(mut self, root: AnyNodeRef<'a>) {
+        if self.enter_node(root).is_traverse() {
+            root.visit_preorder(&mut self);
+        }
+
+        self.leave_node(root);
     }
 
     // Try to skip the subtree if