allow args for the optimize command too

Author: KRRT7

codeflash/cli_cmds/cli.py

@@ -24,10 +24,34 @@ def parse_args() -> Namespace:
     init_actions_parser.set_defaults(func=install_github_actions)
 
     trace_optimize = subparsers.add_parser("optimize", help="Trace and optimize a Python project.")
+
     from codeflash.tracer import main as tracer_main
 
     trace_optimize.set_defaults(func=tracer_main)
 
+    trace_optimize.add_argument(
+        "--max-function-count",
+        type=int,
+        default=100,
+        help="The maximum number of times to trace a single function. More calls to a function will not be traced. Default is 100.",
+    )
+    trace_optimize.add_argument(
+        "--timeout",
+        type=int,
+        help="The maximum time in seconds to trace the entire workflow. Default is indefinite. This is useful while tracing really long workflows, to not wait indefinitely.",
+    )
+    trace_optimize.add_argument(
+        "--output",
+        type=str,
+        default="codeflash.trace",
+        help="The file to save the trace to. Default is codeflash.trace.",
+    )
+    trace_optimize.add_argument(
+        "--config-file-path",
+        type=str,
+        help="The path to the pyproject.toml file which stores the Codeflash config. This is auto-discovered by default.",
+    )
+
     parser.add_argument("--file", help="Try to optimize only this file")
     parser.add_argument("--function", help="Try to optimize only this function within the given file path")
     parser.add_argument(

codeflash/tracer.py

@@ -46,6 +46,7 @@
 from codeflash.verification.verification_utils import get_test_file_path
 
 if TYPE_CHECKING:
+    from argparse import Namespace
     from types import FrameType, TracebackType
 
 
@@ -798,7 +799,7 @@ def runctx(self, cmd: str, global_vars: dict[str, Any], local_vars: dict[str, An
         return self
 
 
-def main() -> ArgumentParser:
+def main(args: Namespace | None = None) -> ArgumentParser:
     parser = ArgumentParser(allow_abbrev=False)
     parser.add_argument("-o", "--outfile", dest="outfile", help="Save trace to <outfile>", default="codeflash.trace")
     parser.add_argument("--only-functions", help="Trace only these functions", nargs="+", default=None)
@@ -824,18 +825,34 @@ def main() -> ArgumentParser:
     )
     parser.add_argument("--trace-only", action="store_true", help="Trace and create replay tests only, don't optimize")
 
-    if not sys.argv[1:]:
-        parser.print_usage()
-        sys.exit(2)
+    if args is not None:
+        parsed_args = args
+        parsed_args.outfile = getattr(args, "output", "codeflash.trace")
+        parsed_args.only_functions = getattr(args, "only_functions", None)
+        parsed_args.max_function_count = getattr(args, "max_function_count", 100)
+        parsed_args.tracer_timeout = getattr(args, "timeout", None)
+        parsed_args.codeflash_config = getattr(args, "config_file_path", None)
+        parsed_args.trace_only = getattr(args, "trace_only", False)
+        parsed_args.module = False
+
+        if getattr(args, "disable", False):
+            console.rule("Codeflash: Tracer disabled by --disable option", style="bold red")
+            return parser
+
+        unknown_args = []
+    else:
+        if not sys.argv[1:]:
+            parser.print_usage()
+            sys.exit(2)
 
-    args, unknown_args = parser.parse_known_args()
-    sys.argv[:] = unknown_args
+        parsed_args, unknown_args = parser.parse_known_args()
+        sys.argv[:] = unknown_args
 
     # The script that we're profiling may chdir, so capture the absolute path
     # to the output file at startup.
-    if args.outfile is not None:
-        args.outfile = Path(args.outfile).resolve()
-    outfile = args.outfile
+    if parsed_args.outfile is not None:
+        parsed_args.outfile = Path(parsed_args.outfile).resolve()
+    outfile = parsed_args.outfile
 
     if len(unknown_args) > 0:
         if args.module: