fixed docstring trimming

Signed-off-by: Ashwin Naren <arihant2math@gmail.com>
RustPython · arihant2math · Jan 31, 2025 · Feb 1, 2025 · Feb 1, 2025 · Feb 1, 2025
commit a130b6cca6ef6a428eff1157df80b9d7cdf81c22
diff --git a/compiler/codegen/src/compile.rs b/compiler/codegen/src/compile.rs
@@ -3357,7 +3357,30 @@ fn split_doc<'a>(
         if let Some(doc) = try_get_constant_string(std::slice::from_ref(&expr.value)) {
             if opts.optimize < 2 {
                 return (
-                    Some(doc.split("\n").map(|x| x.trim_start()).join("\n")),
+                    Some({
+                        let split: Vec<&str> = doc.split('\n').collect();
+                        // Find the amount of whitespace in the first line
+                        let first_text = split.iter().find(|line| !line.trim().is_empty());
+                        if let Some(first_text) = first_text {
+                            let whitespace = first_text
+                                .chars()
+                                .take_while(|c| c.is_whitespace())
+                                .collect::<String>();
+                            split
+                                .iter()
+                                .map(|line| {
+                                    if line.starts_with(&whitespace) {
+                                        line[whitespace.len()..].to_owned()
+                                    } else {
+                                        line.to_string()
+                                    }
+                                })
+                                .collect::<Vec<String>>()
+                                .join("\n")
+                        } else {
+                            doc
+                        }
+                    }),
                     body_rest,
                 );
             } else {