Merge pull request #1031 from codeflash-ai/add-limit-to-codeflash-optimize--m-pytest

Author: KRRT7

codeflash/tracer.py

@@ -57,6 +57,9 @@ def main(args: Namespace | None = None) -> ArgumentParser:
         default=None,
     )
     parser.add_argument("--trace-only", action="store_true", help="Trace and create replay tests only, don't optimize")
+    parser.add_argument(
+        "--limit", type=int, default=None, help="Limit the number of test files to process (for -m pytest mode)"
+    )
 
     if args is not None:
         parsed_args = args
@@ -106,7 +109,7 @@ def main(args: Namespace | None = None) -> ArgumentParser:
             test_paths = []
             replay_test_paths = []
             if parsed_args.module and unknown_args[0] == "pytest":
-                pytest_splits, test_paths = pytest_split(unknown_args[1:])
+                pytest_splits, test_paths = pytest_split(unknown_args[1:], limit=parsed_args.limit)
                 if pytest_splits is None or test_paths is None:
                     console.print(f"❌ Could not find test files in the specified paths: {unknown_args[1:]}")
                     console.print(f"Current working directory: {Path.cwd()}")

codeflash/tracing/pytest_parallelization.py

@@ -7,14 +7,14 @@
 
 
 def pytest_split(
-    arguments: list[str], num_splits: int | None = None
+    arguments: list[str], num_splits: int | None = None, limit: 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.
+        limit: Maximum number of test files to process. If None, processes all files.
 
     Returns:
         List of lists, where each inner list contains test file paths for one group.
@@ -58,6 +58,10 @@ def pytest_split(
     test_files = list(test_files)
     shuffle(test_files)
 
+    # Apply limit if specified
+    if limit is not None and limit > 0:
+        test_files = test_files[:limit]
+
     # 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