fix: normalize module fallback formatting for import merge

codeflash/languages/python/static_analysis/code_extractor.py

@@ -558,7 +558,8 @@ def add_needed_imports_from_module(
         dst_code_fallback = dst_module_code
         parsed_dst_module = None
     else:
-        dst_code_fallback = dst_module_code.code
+        # Keep Module-input fallback formatting aligned with transformed_module.code.lstrip("\n").
+        dst_code_fallback = dst_module_code.code.lstrip("\n")
         parsed_dst_module = dst_module_code
 
     src_module_and_package: ModuleNameAndPackage = calculate_module_and_package(project_root, src_path)

tests/test_add_needed_imports_from_module.py

@@ -527,3 +527,23 @@ def parse():
     return PATTERN.findall("")
 """
     assert result == expected
+
+
+def test_module_input_fallback_strips_leading_newlines() -> None:
+    src_code = """
+def parse():
+    return helper()
+
+def helper():
+    return 1
+"""
+    parsed_module = cst.parse_module(src_code)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        project_root = Path(tmpdir)
+        file_path = project_root / "mod.py"
+        file_path.write_text(src_code)
+
+        result = add_needed_imports_from_module(src_code, parsed_module, file_path, file_path, project_root)
+
+    assert result == src_code.lstrip("\n")