Merge branch 'main' into mihika_pr_631

Saga4 · web-flow · commit 5938547066f7 · 2025-09-12T00:42:00.000+05:30
diff --git a/codeflash/code_utils/code_extractor.py b/codeflash/code_utils/code_extractor.py
@@ -335,12 +335,12 @@ def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> c
         return updated_node
 
 
-def extract_global_statements(source_code: str) -> list[cst.SimpleStatementLine]:
+def extract_global_statements(source_code: str) -> tuple[cst.Module, list[cst.SimpleStatementLine]]:
     """Extract global statements from source code."""
     module = cst.parse_module(source_code)
     collector = GlobalStatementCollector()
     module.visit(collector)
-    return collector.global_statements
+    return module, collector.global_statements
 
 
 def find_last_import_line(target_code: str) -> int:
@@ -373,30 +373,41 @@ def delete___future___aliased_imports(module_code: str) -> str:
 
 
 def add_global_assignments(src_module_code: str, dst_module_code: str) -> str:
-    non_assignment_global_statements = extract_global_statements(src_module_code)
+    src_module, new_added_global_statements = extract_global_statements(src_module_code)
+    dst_module, existing_global_statements = extract_global_statements(dst_module_code)
 
-    # Find the last import line in target
-    last_import_line = find_last_import_line(dst_module_code)
-
-    # Parse the target code
-    target_module = cst.parse_module(dst_module_code)
-
-    # Create transformer to insert non_assignment_global_statements
-    transformer = ImportInserter(non_assignment_global_statements, last_import_line)
-    #
-    # # Apply transformation
-    modified_module = target_module.visit(transformer)
-    dst_module_code = modified_module.code
-
-    # Parse the code
-    original_module = cst.parse_module(dst_module_code)
-    new_module = cst.parse_module(src_module_code)
+    unique_global_statements = []
+    for stmt in new_added_global_statements:
+        if any(
+            stmt is existing_stmt or stmt.deep_equals(existing_stmt) for existing_stmt in existing_global_statements
+        ):
+            continue
+        unique_global_statements.append(stmt)
+
+    mod_dst_code = dst_module_code
+    # Insert unique global statements if any
+    if unique_global_statements:
+        last_import_line = find_last_import_line(dst_module_code)
+        # Reuse already-parsed dst_module
+        transformer = ImportInserter(unique_global_statements, last_import_line)
+        # Use visit inplace, don't parse again
+        modified_module = dst_module.visit(transformer)
+        mod_dst_code = modified_module.code
+        # Parse the code after insertion
+        original_module = cst.parse_module(mod_dst_code)
+    else:
+        # No new statements to insert, reuse already-parsed dst_module
+        original_module = dst_module
 
+    # Parse the src_module_code once only (already done above: src_module)
     # Collect assignments from the new file
     new_collector = GlobalAssignmentCollector()
-    new_module.visit(new_collector)
+    src_module.visit(new_collector)
+    # Only create transformer if there are assignments to insert/transform
+    if not new_collector.assignments:  # nothing to transform
+        return mod_dst_code
 
-    # Transform the original file
+    # Transform the original destination module
     transformer = GlobalAssignmentTransformer(new_collector.assignments, new_collector.assignment_order)
     transformed_module = original_module.visit(transformer)
 
diff --git a/codeflash/code_utils/code_replacer.py b/codeflash/code_utils/code_replacer.py
@@ -412,11 +412,17 @@ def replace_function_definitions_in_module(
     module_abspath: Path,
     preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]],
     project_root_path: Path,
+    should_add_global_assignments: bool = True,  # noqa: FBT001, FBT002
 ) -> bool:
     source_code: str = module_abspath.read_text(encoding="utf8")
     code_to_apply = get_optimized_code_for_module(module_abspath.relative_to(project_root_path), optimized_code)
+
     new_code: str = replace_functions_and_add_imports(
-        add_global_assignments(code_to_apply, source_code),
+        # adding the global assignments before replacing the code, not after
+        # becuase of an "edge case" where the optimized code intoduced a new import and a global assignment using that import
+        # and that import wasn't used before, so it was ignored when calling AddImportsVisitor.add_needed_import inside replace_functions_and_add_imports (because the global assignment wasn't added yet)
+        # this was added at https://github.com/codeflash-ai/codeflash/pull/448
+        add_global_assignments(code_to_apply, source_code) if should_add_global_assignments else source_code,
         function_names,
         code_to_apply,
         module_abspath,
diff --git a/codeflash/code_utils/shell_utils.py b/codeflash/code_utils/shell_utils.py
@@ -15,7 +15,9 @@
     SHELL_RC_EXPORT_PATTERN = re.compile(r"^set CODEFLASH_API_KEY=(cf-.*)$", re.MULTILINE)
     SHELL_RC_EXPORT_PREFIX = "set CODEFLASH_API_KEY="
 else:
-    SHELL_RC_EXPORT_PATTERN = re.compile(r'^(?!#)export CODEFLASH_API_KEY=[\'"]?(cf-[^\s"]+)[\'"]$', re.MULTILINE)
+    SHELL_RC_EXPORT_PATTERN = re.compile(
+        r'^(?!#)export CODEFLASH_API_KEY=(?:"|\')?(cf-[^\s"\']+)(?:"|\')?$', re.MULTILINE
+    )
     SHELL_RC_EXPORT_PREFIX = "export CODEFLASH_API_KEY="
 
 
diff --git a/codeflash/context/unused_definition_remover.py b/codeflash/context/unused_definition_remover.py
@@ -537,6 +537,7 @@ def revert_unused_helper_functions(
                     module_abspath=file_path,
                     preexisting_objects=set(),  # Empty set since we're reverting
                     project_root_path=project_root,
+                    should_add_global_assignments=False,  # since we revert helpers functions after applying the optimization, we know that the file already has global assignments added, otherwise they would be added twice.
                 )
 
                 if reverted_code:
diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py
@@ -53,6 +53,7 @@
 )
 from codeflash.code_utils.env_utils import get_pr_number
 from codeflash.code_utils.formatter import format_code, sort_imports
+from codeflash.code_utils.git_utils import git_root_dir
 from codeflash.code_utils.instrument_existing_tests import inject_profiling_into_existing_test
 from codeflash.code_utils.line_profile_utils import add_decorator_imports
 from codeflash.code_utils.static_analysis import get_first_top_level_function_or_method_ast
@@ -820,7 +821,10 @@ def reformat_code_and_helpers(
         return new_code, new_helper_code
 
     def replace_function_and_helpers_with_optimized_code(
-        self, code_context: CodeOptimizationContext, optimized_code: CodeStringsMarkdown, original_helper_code: str
+        self,
+        code_context: CodeOptimizationContext,
+        optimized_code: CodeStringsMarkdown,
+        original_helper_code: dict[Path, str],
     ) -> bool:
         did_update = False
         read_writable_functions_by_file_path = defaultdict(set)
@@ -1298,11 +1302,13 @@ def process_review(
             "coverage_message": coverage_message,
             "replay_tests": replay_tests,
             "concolic_tests": concolic_tests,
-            "root_dir": self.project_root,
         }
 
         raise_pr = not self.args.no_pr
 
+        if raise_pr or self.args.staging_review:
+            data["root_dir"] = git_root_dir()
+
         if raise_pr and not self.args.staging_review:
             data["git_remote"] = self.args.git_remote
             check_create_pr(**data)
diff --git a/codeflash/tracer.py b/codeflash/tracer.py
@@ -24,6 +24,7 @@
 from codeflash.code_utils.code_utils import get_run_tmp_file
 from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE
 from codeflash.code_utils.config_parser import parse_config_file
+from codeflash.tracing.pytest_parallelization import pytest_split
 
 if TYPE_CHECKING:
     from argparse import Namespace
@@ -86,51 +87,97 @@ def main(args: Namespace | None = None) -> ArgumentParser:
     config, found_config_path = parse_config_file(parsed_args.codeflash_config)
     project_root = project_root_from_module_root(Path(config["module_root"]), found_config_path)
     if len(unknown_args) > 0:
+        args_dict = {
+            "functions": parsed_args.only_functions,
+            "disable": False,
+            "project_root": str(project_root),
+            "max_function_count": parsed_args.max_function_count,
+            "timeout": parsed_args.tracer_timeout,
+            "progname": unknown_args[0],
+            "config": config,
+            "module": parsed_args.module,
+        }
         try:
-            result_pickle_file_path = get_run_tmp_file("tracer_results_file.pkl")
-            args_dict = {
-                "result_pickle_file_path": str(result_pickle_file_path),
-                "output": str(parsed_args.outfile),
-                "functions": parsed_args.only_functions,
-                "disable": False,
-                "project_root": str(project_root),
-                "max_function_count": parsed_args.max_function_count,
-                "timeout": parsed_args.tracer_timeout,
-                "command": " ".join(sys.argv),
-                "progname": unknown_args[0],
-                "config": config,
-                "module": parsed_args.module,
-            }
-
-            subprocess.run(
-                [
-                    SAFE_SYS_EXECUTABLE,
-                    Path(__file__).parent / "tracing" / "tracing_new_process.py",
-                    *sys.argv,
-                    json.dumps(args_dict),
-                ],
-                cwd=Path.cwd(),
-                check=False,
-            )
-            try:
-                with result_pickle_file_path.open(mode="rb") as f:
-                    data = pickle.load(f)
-            except Exception:
-                console.print("❌ Failed to trace. Exiting...")
-                sys.exit(1)
-            finally:
-                result_pickle_file_path.unlink(missing_ok=True)
-
-            replay_test_path = data["replay_test_file_path"]
-            if not parsed_args.trace_only and replay_test_path is not None:
+            pytest_splits = []
+            test_paths = []
+            replay_test_paths = []
+            if parsed_args.module and unknown_args[0] == "pytest":
+                pytest_splits, test_paths = pytest_split(unknown_args[1:])
+
+            if len(pytest_splits) > 1:
+                processes = []
+                test_paths_set = set(test_paths)
+                result_pickle_file_paths = []
+                for i, test_split in enumerate(pytest_splits, start=1):
+                    result_pickle_file_path = get_run_tmp_file(Path(f"tracer_results_file_{i}.pkl"))
+                    result_pickle_file_paths.append(result_pickle_file_path)
+                    args_dict["result_pickle_file_path"] = str(result_pickle_file_path)
+                    outpath = parsed_args.outfile
+                    outpath = outpath.parent / f"{outpath.stem}_{i}{outpath.suffix}"
+                    args_dict["output"] = str(outpath)
+                    updated_sys_argv = []
+                    for elem in sys.argv:
+                        if elem in test_paths_set:
+                            updated_sys_argv.extend(test_split)
+                        else:
+                            updated_sys_argv.append(elem)
+                    args_dict["command"] = " ".join(updated_sys_argv)
+                    processes.append(
+                        subprocess.Popen(
+                            [
+                                SAFE_SYS_EXECUTABLE,
+                                Path(__file__).parent / "tracing" / "tracing_new_process.py",
+                                *updated_sys_argv,
+                                json.dumps(args_dict),
+                            ],
+                            cwd=Path.cwd(),
+                        )
+                    )
+                for process in processes:
+                    process.wait()
+                for result_pickle_file_path in result_pickle_file_paths:
+                    try:
+                        with result_pickle_file_path.open(mode="rb") as f:
+                            data = pickle.load(f)
+                            replay_test_paths.append(str(data["replay_test_file_path"]))
+                    except Exception:
+                        console.print("❌ Failed to trace. Exiting...")
+                        sys.exit(1)
+                    finally:
+                        result_pickle_file_path.unlink(missing_ok=True)
+            else:
+                result_pickle_file_path = get_run_tmp_file(Path("tracer_results_file.pkl"))
+                args_dict["result_pickle_file_path"] = str(result_pickle_file_path)
+                args_dict["output"] = str(parsed_args.outfile)
+                args_dict["command"] = " ".join(sys.argv)
+
+                subprocess.run(
+                    [
+                        SAFE_SYS_EXECUTABLE,
+                        Path(__file__).parent / "tracing" / "tracing_new_process.py",
+                        *sys.argv,
+                        json.dumps(args_dict),
+                    ],
+                    cwd=Path.cwd(),
+                    check=False,
+                )
+                try:
+                    with result_pickle_file_path.open(mode="rb") as f:
+                        data = pickle.load(f)
+                        replay_test_paths.append(str(data["replay_test_file_path"]))
+                except Exception:
+                    console.print("❌ Failed to trace. Exiting...")
+                    sys.exit(1)
+                finally:
+                    result_pickle_file_path.unlink(missing_ok=True)
+            if not parsed_args.trace_only and replay_test_paths:
                 from codeflash.cli_cmds.cli import parse_args, process_pyproject_config
                 from codeflash.cli_cmds.cmd_init import CODEFLASH_LOGO
                 from codeflash.cli_cmds.console import paneled_text
                 from codeflash.telemetry import posthog_cf
                 from codeflash.telemetry.sentry import init_sentry
 
-                sys.argv = ["codeflash", "--replay-test", str(replay_test_path)]
-
+                sys.argv = ["codeflash", "--replay-test", *replay_test_paths]
                 args = parse_args()
                 paneled_text(
                     CODEFLASH_LOGO,
@@ -150,8 +197,8 @@ def main(args: Namespace | None = None) -> ArgumentParser:
                 # Delete the trace file and the replay test file if they exist
                 if outfile:
                     outfile.unlink(missing_ok=True)
-                if replay_test_path:
-                    replay_test_path.unlink(missing_ok=True)
+                for replay_test_path in replay_test_paths:
+                    Path(replay_test_path).unlink(missing_ok=True)
 
         except BrokenPipeError as exc:
             # Prevent "Exception ignored" during interpreter shutdown.
diff --git a/codeflash/tracing/pytest_parallelization.py b/codeflash/tracing/pytest_parallelization.py
@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import os
+from math import ceil
+from pathlib import Path
+from random import shuffle
+
+
+def pytest_split(
+    arguments: list[str], num_splits: int | None = None
+) -> tuple[list[list[str]] | None, list[str] | None]:
+    """Split pytest test files from a directory into N roughly equal groups for parallel execution.
+
+    Args:
+        arguments: List of arguments passed to pytest
+        test_directory: Path to directory containing test files
+        num_splits: Number of groups to split tests into. If None, uses CPU count.
+
+    Returns:
+        List of lists, where each inner list contains test file paths for one group.
+        Returns single list with all tests if number of test files < CPU cores.
+
+    """
+    try:
+        import pytest
+
+        parser = pytest.Parser()
+
+        pytest_args = parser.parse_known_args(arguments)
+        test_paths = getattr(pytest_args, "file_or_dir", None)
+        if not test_paths:
+            return None, None
+
+    except ImportError:
+        return None, None
+    test_files = set()
+
+    # Find all test_*.py files recursively in the directory
+    for test_path in test_paths:
+        _test_path = Path(test_path)
+        if not _test_path.exists():
+            return None, None
+        if _test_path.is_dir():
+            # Find all test files matching the pattern test_*.py
+            test_files.update(map(str, _test_path.rglob("test_*.py")))
+            test_files.update(map(str, _test_path.rglob("*_test.py")))
+        elif _test_path.is_file():
+            test_files.add(str(_test_path))
+
+    if not test_files:
+        return [[]], None
+
+    # Determine number of splits
+    if num_splits is None:
+        num_splits = os.cpu_count() or 4
+
+    # randomize to increase chances of all splits being balanced
+    test_files = list(test_files)
+    shuffle(test_files)
+
+    # Ensure each split has at least 4 test files
+    # If we have fewer test files than 4 * num_splits, reduce num_splits
+    max_possible_splits = len(test_files) // 4
+    if max_possible_splits == 0:
+        return test_files, test_paths
+
+    num_splits = min(num_splits, max_possible_splits)
+
+    # Calculate chunk size (round up to ensure all files are included)
+    total_files = len(test_files)
+    chunk_size = ceil(total_files / num_splits)
+
+    # Initialize result groups
+    result_groups = [[] for _ in range(num_splits)]
+
+    # Distribute files across groups
+    for i, test_file in enumerate(test_files):
+        group_index = i // chunk_size
+        # Ensure we don't exceed the number of groups (edge case handling)
+        if group_index >= num_splits:
+            group_index = num_splits - 1
+        result_groups[group_index].append(test_file)
+
+    return result_groups, test_paths
diff --git a/codeflash/version.py b/codeflash/version.py
@@ -1,2 +1,2 @@
 # These version placeholders will be replaced by uv-dynamic-versioning during build.
-__version__ = "0.16.6"
+__version__ = "0.16.7"
diff --git a/pyproject.toml b/pyproject.toml
diff --git a/tests/test_code_replacement.py b/tests/test_code_replacement.py
diff --git a/tests/test_shell_utils.py b/tests/test_shell_utils.py
diff --git a/uv.lock b/uv.lock

Original file line number	Diff line number	Diff line change
`@@ -537,6 +537,7 @@ def revert_unused_helper_functions(`
`537`	`537`	`module_abspath=file_path,`
`538`	`538`	`preexisting_objects=set(), # Empty set since we're reverting`
`539`	`539`	`project_root_path=project_root,`
	`540`	`+ should_add_global_assignments=False, # since we revert helpers functions after applying the optimization, we know that the file already has global assignments added, otherwise they would be added twice.`
`540`	`541`	`)`
`541`	`542`
`542`	`543`	`if reverted_code:`
Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,2 @@`
`1`	`1`	`# These version placeholders will be replaced by uv-dynamic-versioning during build.`
`2`		`-__version__ = "0.16.6"`
	`2`	`+__version__ = "0.16.7"`