fix(docstring): retain doctest arrow prefix

crates/pyrefly_python/src/docstring.rs

@@ -5,7 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-use std::cmp::min;
 use std::collections::HashMap;
 
 use ruff_python_ast::Expr;
@@ -66,24 +65,30 @@ impl Docstring {
                 if i == 0 {
                     line.to_owned()
                 } else {
-                    let trimmed = &line[min(min_indent, line.len())..];
-                    let mut without_blockquote = trimmed;
-                    while let Some(rest) = without_blockquote.strip_prefix('>') {
-                        without_blockquote = rest.strip_prefix(' ').unwrap_or(rest);
+                    let trimmed = &line[min_indent.min(line.len())..];
+                    let mut content = trimmed;
+
+                    // Handle potential leading blockquote (`> `) for non-doctest lines
+                    let is_doctest_prompt = {
+                        let t = trimmed.trim_start();
+                        t.starts_with(">>>") && t.as_bytes().get(3).is_none_or(|b| *b != b'>')
+                    };
+                    if !is_doctest_prompt {
+                        while let Some(rest) = content.strip_prefix('>') {
+                            content = rest.strip_prefix(' ').unwrap_or(rest);
+                        }
                     }
+
                     // Replace remaining leading spaces with   or they might be ignored in markdown parsers
-                    let leading_spaces = without_blockquote
-                        .bytes()
-                        .take_while(|&c| c == b' ')
-                        .count();
+                    let leading_spaces = content.bytes().take_while(|&c| c == b' ').count();
                     if leading_spaces > 0 {
                         format!(
                             "{}{}",
                             " ".repeat(leading_spaces),
-                            &without_blockquote[leading_spaces..]
+                            &content[leading_spaces..]
                         )
                     } else {
-                        without_blockquote.to_owned()
+                        content.to_owned()
                     }
                 }
             })
@@ -454,6 +459,15 @@ mod tests {
             "hello  \nworld  \ntest"
         );
     }
+
+    #[test]
+    fn test_docstring_preserves_doctest_prompt() {
+        assert_eq!(
+            Docstring::clean("\"\"\"Example\n>>> foo()\"\"\"").as_str(),
+            "Example  \n>>> foo()"
+        );
+    }
+
     #[test]
     fn test_parse_sphinx_param_docs() {
         let doc = r#"

pyrefly/lib/test/lsp/hover_docstring.rs

@@ -379,30 +379,3 @@ Docstring Result: `Test docstring`
         report.trim(),
     );
 }
-
-#[test]
-fn test_removes_block_quote_symbols_at_start_of_line() {
-    let code = r#"
-def fun() -> None:
-    """
-    >>> d = {"col1": [1, 2], "col2": [3, 4]}
-    """
-    pass
-
-f = fun()
-#   ^
-"#;
-    let report = get_batched_lsp_operations_report(&[("main", code)], test_report_factory(code));
-    assert_eq!(
-        r#"
-# main.py
-8 | f = fun()
-        ^
-Docstring Result: `  
-d = {"col1": [1, 2], "col2": [3, 4]}  
-`
-"#
-        .trim(),
-        report.trim(),
-    );
-}