pyo3_stub_gen/generate/
docstring.rs

1use std::fmt;
2
3/// Normalize a docstring by trimming outer whitespace and dedenting.
4///
5/// Implements Python's inspect.cleandoc() behavior:
6/// 1. Trim leading/trailing whitespace from the entire string
7/// 2. Find minimum indentation of non-empty lines (skip first line)
8/// 3. Remove that indentation from all lines
9///
10/// # Examples
11/// ```
12/// # use pyo3_stub_gen::generate::normalize_docstring;
13/// let doc = r#"
14///     First line
15///     Second line
16///         Indented line
17/// "#;
18/// let normalized = normalize_docstring(doc);
19/// assert_eq!(normalized, "First line\nSecond line\n    Indented line");
20/// ```
21pub fn normalize_docstring(doc: &str) -> String {
22    let doc = doc.trim();
23    if doc.is_empty() {
24        return String::new();
25    }
26
27    let lines: Vec<&str> = doc.lines().collect();
28
29    // Find minimum indentation of non-empty lines (skip first line)
30    let min_indent = lines
31        .iter()
32        .skip(1)
33        .filter(|line| !line.trim().is_empty())
34        .map(|line| line.chars().take_while(|c| c.is_whitespace()).count())
35        .min()
36        .unwrap_or(0);
37
38    // Build normalized lines with dedenting applied
39    let normalized_lines: Vec<String> = lines
40        .iter()
41        .enumerate()
42        .map(|(i, line)| {
43            if i == 0 {
44                // First line: use as-is (already trimmed by outer trim())
45                line.to_string()
46            } else if line.trim().is_empty() {
47                // Empty line: keep it but remove whitespace
48                String::new()
49            } else {
50                // Other lines: remove common indentation
51                if line.len() >= min_indent {
52                    line[min_indent..].to_string()
53                } else {
54                    line.trim_start().to_string()
55                }
56            }
57        })
58        .collect();
59
60    normalized_lines.join("\n")
61}
62
63pub fn write_docstring(f: &mut impl fmt::Write, doc: &str, indent: &str) -> fmt::Result {
64    // Docstrings should already be normalized, but trim again for safety
65    let doc = doc.trim();
66    if !doc.is_empty() {
67        let needs_escaping = doc.contains("\"\"\"") || doc.contains(['\0', '\r']);
68        let prefix = if needs_escaping { "" } else { "r" };
69        writeln!(f, "{indent}{prefix}\"\"\"")?;
70        for line in doc.lines() {
71            write!(f, "{indent}")?;
72            if needs_escaping {
73                for ch in line.chars() {
74                    match ch {
75                        '\\' => f.write_str("\\\\")?,
76                        '"' => f.write_str("\\\"")?,
77                        '\0' => f.write_str("\\x00")?,
78                        '\r' => f.write_str("\\r")?,
79                        _ => f.write_char(ch)?,
80                    }
81                }
82            } else {
83                f.write_str(line)?;
84            }
85            writeln!(f)?;
86        }
87        writeln!(f, r#"{indent}""""#)?;
88    }
89    Ok(())
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95    use crate::generate::Module;
96    use pyo3::{prelude::*, types::PyModule};
97    use std::ffi::CString;
98    use test_case::test_case;
99
100    #[test_case("First line\n\n    Indented line"; "multiline")]
101    #[test_case(r#"Example: """triple quoted""" and '''single quoted'''."#; "triple_quotes")]
102    #[test_case(r#"Quoted: \"\"\" and \"""."#; "escaped_quotes")]
103    #[test_case(r#"Paths: C:\new\tools\ and escapes: \n \t \x00 \u1234."#; "literal_backslashes")]
104    #[test_case(r#"Code: """C:\new\tools""" and \n \u1234."#; "quoted_literal_backslashes")]
105    #[test_case("Line ends with \\\nnext line"; "backslash_before_newline")]
106    #[test_case("Trailing backslash \\"; "trailing_backslash")]
107    #[test_case("Triple quotes \"\"\" and trailing backslash \\"; "escaped_trailing_backslash")]
108    #[test_case("Control: \0 and \r inside"; "control_characters")]
109    #[test_case("Unicode: 日本語 🦀 with \"\"\" quotes"; "unicode")]
110    fn docstrings_round_trip_through_python(doc: &str) {
111        let mut source = Module {
112            doc: doc.into(),
113            ..Default::default()
114        }
115        .format_init_py();
116        source.push_str("\ndef documented():\n");
117        write_docstring(&mut source, doc, "    ").unwrap();
118        source.push_str("    pass\n");
119
120        Python::initialize();
121        Python::attach(|py| {
122            let source = CString::new(source).unwrap();
123            let module = PyModule::new(py, "docstring_test").unwrap();
124            py.run(&source, Some(&module.dict()), None).unwrap();
125            let inspect = py.import("inspect").unwrap();
126            for object in [
127                module.clone().into_any(),
128                module.getattr("documented").unwrap(),
129            ] {
130                let actual: String = inspect
131                    .call_method1("getdoc", (object,))
132                    .unwrap()
133                    .extract()
134                    .unwrap();
135                assert_eq!(actual, doc);
136            }
137        });
138    }
139}
pyo3_stub_gen/generate/docstring.rs

pyo3_stub_gen/generate/
docstring.rs
