From 3444babf58ee522b32b1ec2fef004ffde7df0a4c Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 3 Apr 2026 11:15:18 -0500 Subject: [PATCH 1/2] add --inject flag to codeflash compare When benchmarking already-merged optimizations, the benchmark file often doesn't exist at either the base or head ref. The --inject flag copies specified files/directories from the working tree into both worktrees before benchmark discovery and execution, eliminating the need to cherry-pick benchmark commits onto temporary branches. Usage: codeflash compare --inject tests/benchmarks/test_bench.py --- codeflash/benchmarking/compare.py | 19 +++++++++++++++++++ codeflash/cli_cmds/cli.py | 6 ++++++ codeflash/cli_cmds/cmd_compare.py | 1 + 3 files changed, 26 insertions(+) diff --git a/codeflash/benchmarking/compare.py b/codeflash/benchmarking/compare.py index 9ce4db01b..068c4fa99 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 as _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..16fa8ea10 100644 --- a/codeflash/cli_cmds/cmd_compare.py +++ b/codeflash/cli_cmds/cmd_compare.py @@ -108,6 +108,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: From a5a475d0b425fc69bff037ac3004783e97563a8c Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 3 Apr 2026 11:21:05 -0500 Subject: [PATCH 2/2] address review: fix _shutil alias, warn on --inject with --script - Replace `import shutil as _shutil` with plain `import shutil` to match the existing style in the same function - Warn when --inject is used with --script mode (unsupported combo) instead of silently dropping the flag --- codeflash/benchmarking/compare.py | 6 +++--- codeflash/cli_cmds/cmd_compare.py | 3 +++ 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/codeflash/benchmarking/compare.py b/codeflash/benchmarking/compare.py index 068c4fa99..237753cb8 100644 --- a/codeflash/benchmarking/compare.py +++ b/codeflash/benchmarking/compare.py @@ -347,7 +347,7 @@ def build_panel(current_step: int) -> Panel: # Inject files from working tree into both worktrees if inject_paths: - import shutil as _shutil + import shutil for path_str in inject_paths: src = repo_root / path_str @@ -358,9 +358,9 @@ def build_panel(current_step: int) -> Panel: dst = wt / path_str dst.parent.mkdir(parents=True, exist_ok=True) if src.is_dir(): - _shutil.copytree(src, dst, dirs_exist_ok=True) + shutil.copytree(src, dst, dirs_exist_ok=True) elif src.is_file(): - _shutil.copy2(src, dst) + shutil.copy2(src, dst) step += 1 live.update(build_panel(step)) diff --git a/codeflash/cli_cmds/cmd_compare.py b/codeflash/cli_cmds/cmd_compare.py index 16fa8ea10..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")