Better formatting for >>> in docstrings, fixes #49, fixes #163

Author: davidhalter

crates/zuban_python/src/name.rs

@@ -340,7 +340,7 @@ impl<'db, 'x> Name<'db, 'x> {
             debug_assert!(!name.file().is_stub());
             return name.documentation();
         }
-        result
+        docstr_to_markdown(result)
     }
 
     pub fn class_symbols(&self) -> Option<impl ExactSizeIterator<Item = NameSymbol<'db>>> {
@@ -564,3 +564,32 @@ impl<'db> NameSymbol<'db> {
         })
     }
 }
+
+fn docstr_to_markdown(s: Cow<str>) -> Cow<str> {
+    if !s.contains(">>>") {
+        return s;
+    }
+    let mut out_parts = vec![];
+    let mut in_block = false;
+
+    for line in s.split_inclusive('\n') {
+        if line.starts_with(">>>") {
+            if !in_block {
+                out_parts.push("```python\n");
+                in_block = true;
+            }
+        } else if in_block && line.trim().is_empty() {
+            out_parts.push("```\n");
+            in_block = false;
+        }
+        out_parts.push(line);
+    }
+
+    if in_block {
+        if !out_parts.last().unwrap().ends_with('\n') {
+            out_parts.push("\n");
+        }
+        out_parts.push("```\n");
+    }
+    Cow::Owned(out_parts.join(""))
+}

crates/zuban_python/tests/mypylike/tests/documentation.test

@@ -235,3 +235,23 @@ D
 __main__.py:12:documentation -> "```python\n(class) A(a, *, a2: str) -> A\n```\n---\na doc"
 __main__.py:14:documentation -> "```python\n(class) B() -> B\n```"
 __main__.py:16:documentation -> "```python\n(class) D(x: int, y: str = ...) -> D\n```"
+
+[case documentation_doctests_formatting]
+def f() -> None:
+    """
+    asdf
+    >>> bla
+    foo
+    >>> 1
+    bar
+
+    x
+    
+    >>> hello
+    """
+
+#? documentation
+f
+
+[out]
+__main__.py:15:documentation -> "```python\n(function) def f() -> None\n```\n---\nasdf\n```python\n>>> bla\nfoo\n>>> 1\nbar\n```\n\nx\n\n```python\n>>> hello\n```\n"