From bf9adf2673029058746745e659d2f4f86523366d Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Tue, 17 Mar 2026 18:16:54 -0600 Subject: [PATCH] fix: normalize module fallback formatting for import merge --- .../python/static_analysis/code_extractor.py | 3 ++- tests/test_add_needed_imports_from_module.py | 20 +++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/codeflash/languages/python/static_analysis/code_extractor.py b/codeflash/languages/python/static_analysis/code_extractor.py index 2570c1b90..6bd81bb74 100644 --- a/codeflash/languages/python/static_analysis/code_extractor.py +++ b/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) diff --git a/tests/test_add_needed_imports_from_module.py b/tests/test_add_needed_imports_from_module.py index 345b966dc..198058b28 100644 --- a/tests/test_add_needed_imports_from_module.py +++ b/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")