fixed docstring trimming

Signed-off-by: Ashwin Naren <arihant2math@gmail.com>
RustPython · youknowone · Feb 13, 2025 · Jan 31, 2025 · Feb 1, 2025 · Feb 1, 2025
commit 548e1f2e620b403bdc7e7e9d180eb7b902402fd0
diff --git a/compiler/codegen/src/compile.rs b/compiler/codegen/src/compile.rs
@@ -3376,7 +3376,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 {