Support formatting pycon markdown code blocks (#23112)

Fix #23078

Author: amyreese

crates/ruff_markdown/src/lib.rs

@@ -5,7 +5,7 @@ use ruff_python_ast::PySourceType;
 use ruff_python_formatter::format_module_source;
 use ruff_python_trivia::textwrap::{dedent, indent};
 use ruff_source_file::{Line, UniversalNewlines};
-use ruff_text_size::{TextRange, TextSize};
+use ruff_text_size::{TextLen, TextRange, TextSize};
 use ruff_workspace::FormatterSettings;
 
 #[derive(Debug, PartialEq, Eq)]
@@ -54,7 +54,7 @@ pub fn format_code_blocks(
     let mut state = MarkdownState::On;
     let mut changed = false;
     let mut formatted = String::with_capacity(source.len());
-    let mut last_match = TextSize::new(0);
+    let mut last_match = TextSize::ZERO;
 
     let mut lines = source.universal_newlines().peekable();
     while let Some(line) = lines.next() {
@@ -80,46 +80,52 @@ pub fn format_code_blocks(
                     continue;
                 };
 
+                if closing_fence != opening_fence {
+                    continue;
+                }
+
                 // Found the matching end of the code block
-                if closing_fence == opening_fence {
-                    let language = language.to_ascii_lowercase();
-                    if state == MarkdownState::On
-                        && matches!(
-                            language.as_str(),
-                            "python" | "py" | "python3" | "py3" | "pyi"
-                        )
-                    {
-                        // Maybe python, try formatting it
-                        let end = code_line.start();
-                        let unformatted_code = dedent(&source[TextRange::new(start, end)]);
-
-                        let py_source_type = match settings.extension.get_extension(&language) {
-                            None => PySourceType::from_extension(&language),
-                            Some(language) => PySourceType::from(language),
-                        };
+                if state != MarkdownState::On {
+                    break;
+                }
+
+                // Maybe python, try formatting it
+                let language = language.to_ascii_lowercase();
+                let py_source_type = match settings.extension.get_extension(&language) {
+                    None => PySourceType::from_extension(&language),
+                    Some(language) => PySourceType::from(language),
+                };
+
+                let end = code_line.start();
+                let unformatted_code = dedent(&source[TextRange::new(start, end)]);
+
+                let formatted_code = match language.as_str() {
+                    "python" | "py" | "python3" | "py3" | "pyi" => {
                         let options =
                             settings.to_format_options(py_source_type, &unformatted_code, path);
-
                         // Using `Printed::into_code` requires adding `ruff_formatter` as a direct
                         // dependency, and I suspect that Rust can optimize the closure away regardless.
                         #[expect(clippy::redundant_closure_for_method_calls)]
-                        let formatted_code = format_module_source(&unformatted_code, options)
-                            .map(|formatted| formatted.into_code());
-
-                        // Formatting produced changes
-                        if let Ok(formatted_code) = formatted_code
-                            && (formatted_code.len() != unformatted_code.len()
-                                || formatted_code != *unformatted_code)
-                        {
-                            formatted.push_str(&source[TextRange::new(last_match, start)]);
-                            let formatted_code = indent(&formatted_code, code_indent);
-                            formatted.push_str(&formatted_code);
-                            last_match = end;
-                            changed = true;
-                        }
+                        format_module_source(&unformatted_code, options)
+                            .map(|formatted| formatted.into_code())
+                            .ok()
                     }
-                    break;
+                    "pycon" => format_pycon_block(&unformatted_code, path, settings),
+                    _ => None,
+                };
+
+                // Formatting produced changes
+                if let Some(formatted_code) = formatted_code
+                    && (formatted_code.len() != unformatted_code.len()
+                        || formatted_code != *unformatted_code)
+                {
+                    formatted.push_str(&source[TextRange::new(last_match, start)]);
+                    let formatted_code = indent(&formatted_code, code_indent);
+                    formatted.push_str(&formatted_code);
+                    last_match = end;
+                    changed = true;
                 }
+                break;
             }
         }
     }
@@ -132,6 +138,73 @@ pub fn format_code_blocks(
     }
 }
 
+fn format_pycon_block(
+    source: &str,
+    path: Option<&Path>,
+    settings: &FormatterSettings,
+) -> Option<String> {
+    static FIRST_LINE: &str = ">>> ";
+    static CONTINUATION: &str = "... ";
+    static CONTINUATION_BLANK: &str = "...";
+
+    let offset = FIRST_LINE.text_len();
+    let mut changed = false;
+    let mut result = String::with_capacity(source.len());
+    let mut unformatted = String::with_capacity(source.len());
+    let mut last_match = TextSize::new(0);
+    let mut lines = source.universal_newlines().peekable();
+
+    while let Some(line) = lines.next() {
+        unformatted.clear();
+        if line.starts_with(FIRST_LINE) {
+            let start = line.start();
+            let mut end = line.full_end();
+            unformatted.push_str(&source[TextRange::new(line.start() + offset, line.full_end())]);
+            while let Some(next_line) = lines.next_if(|line| line.starts_with(CONTINUATION_BLANK)) {
+                end = next_line.full_end();
+                let start = if next_line.trim_end() == CONTINUATION_BLANK {
+                    next_line.end()
+                } else {
+                    next_line.start() + offset
+                };
+                unformatted.push_str(&source[TextRange::new(start, end)]);
+            }
+            let options = settings.to_format_options(PySourceType::Python, &unformatted, path);
+            // Using `Printed::into_code` requires adding `ruff_formatter` as a direct
+            // dependency, and I suspect that Rust can optimize the closure away regardless.
+            #[expect(clippy::redundant_closure_for_method_calls)]
+            let Ok(formatted) =
+                format_module_source(&unformatted, options).map(|formatted| formatted.into_code())
+            else {
+                continue;
+            };
+
+            if formatted.len() != unformatted.len() || formatted != unformatted {
+                result.push_str(&source[TextRange::new(last_match, start)]);
+                for (idx, line) in formatted.universal_newlines().enumerate() {
+                    result.push_str(if idx == 0 {
+                        FIRST_LINE
+                    } else if line.is_empty() {
+                        CONTINUATION_BLANK
+                    } else {
+                        CONTINUATION
+                    });
+                    result.push_str(&formatted[line.full_range()]);
+                }
+                last_match = end;
+                changed = true;
+            }
+        }
+    }
+
+    if changed {
+        result.push_str(&source[last_match.to_usize()..]);
+        Some(result)
+    } else {
+        None
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use insta::assert_snapshot;
@@ -431,4 +504,39 @@ def bar(): ...
         ~~~
         "#);
     }
+
+    #[test]
+    fn format_code_blocks_python_console() {
+        let code = r#"
+```pycon
+>>> print( 'hello there' )
+hello there
+>>> def foo(): pass
+>>> def bar():
+...   print( 'thing1', "thing2", )
+...
+... bar()
+...
+thing1 thing2
+```
+        "#;
+        assert_snapshot!(format_code_blocks(code, None, &FormatterSettings::default()), @r#"
+
+        ```pycon
+        >>> print("hello there")
+        hello there
+        >>> def foo():
+        ...     pass
+        >>> def bar():
+        ...     print(
+        ...         "thing1",
+        ...         "thing2",
+        ...     )
+        ...
+        ...
+        ... bar()
+        thing1 thing2
+        ```
+        "#);
+    }
 }