diff --git a/codeflash/benchmarking/compare.py b/codeflash/benchmarking/compare.py index 9ce4db01b..237753cb8 100644 --- a/codeflash/benchmarking/compare.py +++ b/codeflash/benchmarking/compare.py @@ -211,6 +211,7 @@ def compare_branches( functions: Optional[dict[Path, list[FunctionToOptimize]]] = None, timeout: int = 600, memory: bool = False, + inject_paths: Optional[list[str]] = None, ) -> CompareResult: """Compare benchmark performance between two git refs. @@ -343,6 +344,24 @@ def build_panel(current_step: int) -> Panel: head_sha = repo.commit(head_ref).hexsha repo.git.worktree("add", str(base_worktree), base_sha) repo.git.worktree("add", str(head_worktree), head_sha) + + # Inject files from working tree into both worktrees + if inject_paths: + import shutil + + for path_str in inject_paths: + src = repo_root / path_str + if not src.exists(): + logger.warning("Inject path does not exist: %s", src) + continue + for wt in [base_worktree, head_worktree]: + dst = wt / path_str + dst.parent.mkdir(parents=True, exist_ok=True) + if src.is_dir(): + shutil.copytree(src, dst, dirs_exist_ok=True) + elif src.is_file(): + shutil.copy2(src, dst) + step += 1 live.update(build_panel(step)) diff --git a/codeflash/cli_cmds/cli.py b/codeflash/cli_cmds/cli.py index cf5ca7bdd..954a1165d 100644 --- a/codeflash/cli_cmds/cli.py +++ b/codeflash/cli_cmds/cli.py @@ -403,6 +403,12 @@ def _build_parser() -> ArgumentParser: help="Relative path to JSON results file produced by --script (required with --script)", ) compare_parser.add_argument("--config-file", type=str, dest="config_file", help="Path to pyproject.toml") + compare_parser.add_argument( + "--inject", + nargs="+", + default=None, + help="Files or directories to copy into both worktrees before benchmarking. Paths are relative to repo root.", + ) trace_optimize = subparsers.add_parser("optimize", help="Trace and optimize your project.") diff --git a/codeflash/cli_cmds/cmd_compare.py b/codeflash/cli_cmds/cmd_compare.py index 898af5679..87d659fdb 100644 --- a/codeflash/cli_cmds/cmd_compare.py +++ b/codeflash/cli_cmds/cmd_compare.py @@ -40,6 +40,9 @@ def run_compare(args: Namespace) -> None: # Script mode: run an arbitrary benchmark command on each worktree (no codeflash config needed) script_cmd = getattr(args, "script", None) if script_cmd: + if getattr(args, "inject", None): + logger.warning("--inject is not supported in --script mode and will be ignored") + script_output = getattr(args, "script_output", None) if not script_output: logger.error("--script-output is required when using --script") @@ -108,6 +111,7 @@ def run_compare(args: Namespace) -> None: functions=functions, timeout=args.timeout, memory=getattr(args, "memory", False), + inject_paths=getattr(args, "inject", None), ) if not result.base_stats and not result.head_stats: