ruff_python_formatting: add doctest state handling

This commit adds a small but central component to code snippet
formatting in docstrings: it specifically implements the state
transitions needed to recognize and collect code snippets from
doctests. This means looking for PS1 and PS2 prompts and extracting
the code portion of each line.

This also introduces a "code example add action" which we will
use in a subsequent commit to control the higher level docstring
line printer.

Author: BurntSushi

crates/ruff_python_formatter/src/expression/string.rs

@@ -1071,6 +1071,83 @@ impl<'src> DocstringLine<'src> {
         }
     }
 }
+
+/// A single code example extracted from a docstring.
+///
+/// This represents an intermediate state from when the code example was first
+/// found all the way up until the point at which the code example has finished
+/// and is reformatted.
+///
+/// Its default state is "empty." That is, that no code example is currently
+/// being collected.
+#[derive(Debug, Default)]
+struct CodeExample<'src> {
+    /// The kind of code example being collected, or `None` if no code example
+    /// has been observed.
+    kind: Option<CodeExampleKind>,
+    /// The lines that have been seen so far that make up the code example.
+    lines: Vec<CodeExampleLine<'src>>,
+}
+
+impl<'src> CodeExample<'src> {
+    /// Attempt to add an original line from a docstring to this code example.
+    ///
+    /// Based on the line and the internal state of whether a code example is
+    /// currently being collected or not, this will return an "action" for
+    /// the caller to perform. The typical case is a "print" action, which
+    /// instructs the caller to just print the line as though it were not part
+    /// of a code snippet.
+    fn add(&mut self, original: DocstringLine<'src>) -> CodeExampleAddAction<'src> {
+        match self.kind.take() {
+            // There's no existing code example being built, so we look for
+            // the start of one or otherwise tell the caller we couldn't find
+            // anything.
+            None => match self.add_start(original) {
+                None => CodeExampleAddAction::Kept,
+                Some(original) => CodeExampleAddAction::Print { original },
+            },
+            Some(CodeExampleKind::Doctest(doctest)) => {
+                if let Some(code) = doctest_find_ps2_prompt(&doctest.indent, &original.line) {
+                    let code = code.to_string();
+                    self.lines.push(CodeExampleLine { original, code });
+                    // Stay with the doctest kind while we accumulate all
+                    // PS2 prompts.
+                    self.kind = Some(CodeExampleKind::Doctest(doctest));
+                    return CodeExampleAddAction::Kept;
+                }
+                let code = std::mem::take(&mut self.lines);
+                let original = self.add_start(original);
+                CodeExampleAddAction::Format {
+                    code,
+                    kind: CodeExampleKind::Doctest(doctest),
+                    original,
+                }
+            }
+        }
+    }
+
+    /// Looks for the start of a code example. If one was found, then the given
+    /// line is kept and added as part of the code example. Otherwise, the line
+    /// is returned unchanged and no code example was found.
+    ///
+    /// # Panics
+    ///
+    /// This panics when the existing code-example is any non-None value. That
+    /// is, this routine assumes that there is no ongoing code example being
+    /// collected and looks for the beginning of another code example.
+    fn add_start(&mut self, original: DocstringLine<'src>) -> Option<DocstringLine<'src>> {
+        assert_eq!(None, self.kind, "expected no existing code example");
+        if let Some((indent, code)) = doctest_find_ps1_prompt(&original.line) {
+            let indent = indent.to_string();
+            let code = code.to_string();
+            self.lines.push(CodeExampleLine { original, code });
+            self.kind = Some(CodeExampleKind::Doctest(CodeExampleDoctest { indent }));
+            return None;
+        }
+        Some(original)
+    }
+}
+
 /// The kind of code example observed in a docstring.
 #[derive(Clone, Debug, Eq, PartialEq)]
 enum CodeExampleKind {
@@ -1117,6 +1194,97 @@ struct CodeExampleLine<'src> {
     /// The code extracted from the line.
     code: String,
 }
+
+/// An action that a caller should perform after attempting to add a line from
+/// a docstring to a code example.
+///
+/// Callers are expected to add every line from a docstring to a code example,
+/// and the state of the code example (and the line itself) will determine
+/// how the caller should react.
+#[derive(Debug)]
+enum CodeExampleAddAction<'src> {
+    /// The line added was ignored by `CodeExample` and the caller should print
+    /// it to the formatter as-is.
+    ///
+    /// This is the common case. That is, most lines in most docstrings are not
+    /// part of a code example.
+    Print { original: DocstringLine<'src> },
+    /// The line added was kept by `CodeExample` as part of a new or existing
+    /// code example.
+    ///
+    /// When this occurs, callers should not try to format the line and instead
+    /// move on to the next line.
+    Kept,
+    /// The line added indicated that the code example is finished and should
+    /// be formatted and printed. The line added is not treated as part of
+    /// the code example. If the line added indicated the start of another
+    /// code example, then is won't be returned to the caller here. Otherwise,
+    /// callers should pass it through to the formatter as-is.
+    Format {
+        /// The kind of code example that was found.
+        kind: CodeExampleKind,
+        /// The Python code that should be formatted, indented and printed.
+        ///
+        /// This is guaranteed to be non-empty.
+        code: Vec<CodeExampleLine<'src>>,
+        /// When set, the line is considered not part of any code example
+        /// and should be formatted as if the `Ignore` action were returned.
+        /// Otherwise, if there is no line, then either one does not exist
+        /// or it is part of another code example and should be treated as a
+        /// `Kept` action.
+        original: Option<DocstringLine<'src>>,
+    },
+}
+
+/// Looks for a valid doctest PS1 prompt in the line given.
+///
+/// If one was found, then the indentation prior to the prompt is returned
+/// along with the code portion of the line.
+fn doctest_find_ps1_prompt(line: &str) -> Option<(&str, &str)> {
+    let trim_start = line.trim_start();
+    // Prompts must be followed by an ASCII space character[1].
+    //
+    // [1]: https://github.com/python/cpython/blob/0ff6368519ed7542ad8b443de01108690102420a/Lib/doctest.py#L809-L812
+    let code = trim_start.strip_prefix(">>> ")?;
+    let indent_len = line
+        .len()
+        .checked_sub(trim_start.len())
+        .expect("suffix is <= original");
+    let indent = &line[..indent_len];
+    Some((indent, code))
+}
+
+/// Looks for a valid doctest PS2 prompt in the line given.
+///
+/// If one is found, then the code portion of the line following the PS2 prompt
+/// is returned.
+///
+/// Callers must provide a string containing the original indentation of the
+/// PS1 prompt that started the doctest containing the potential PS2 prompt
+/// in the line given. If the line contains a PS2 prompt, its indentation must
+/// match the indentation used for the corresponding PS1 prompt (otherwise
+/// `None` will be returned).
+fn doctest_find_ps2_prompt<'src>(ps1_indent: &str, line: &'src str) -> Option<&'src str> {
+    let (ps2_indent, ps2_after) = line.split_once("...")?;
+    // PS2 prompts must have the same indentation as their
+    // corresponding PS1 prompt.[1] While the 'doctest' Python
+    // module will error in this case, we just treat this line as a
+    // non-doctest line.
+    //
+    // [1]: https://github.com/python/cpython/blob/0ff6368519ed7542ad8b443de01108690102420a/Lib/doctest.py#L733
+    if ps1_indent != ps2_indent {
+        return None;
+    }
+    // PS2 prompts must be followed by an ASCII space character unless
+    // it's an otherwise empty line[1].
+    //
+    // [1]: https://github.com/python/cpython/blob/0ff6368519ed7542ad8b443de01108690102420a/Lib/doctest.py#L809-L812
+    match ps2_after.strip_prefix(' ') {
+        None if ps2_after.is_empty() => Some(""),
+        None => None,
+        Some(code) => Some(code),
+    }
+}
 }
 
 /// If the last line of the docstring is `content" """` or `content\ """`, we need a chaperone space