@@ -123,8 +123,8 @@ def __init__(
123
123
self .function_count = defaultdict (int )
124
124
self .current_file_path = Path (__file__ ).resolve ()
125
125
self .ignored_qualified_functions = {
126
- f"{ self .current_file_path } :Tracer: __exit__" ,
127
- f"{ self .current_file_path } :Tracer: __enter__" ,
126
+ f"{ self .current_file_path } :Tracer. __exit__" ,
127
+ f"{ self .current_file_path } :Tracer. __enter__" ,
128
128
}
129
129
self .max_function_count = max_function_count
130
130
self .config , found_config_path = parse_config_file (config_file_path )
@@ -133,6 +133,7 @@ def __init__(
133
133
self .ignored_functions = {"<listcomp>" , "<genexpr>" , "<dictcomp>" , "<setcomp>" , "<lambda>" , "<module>" }
134
134
135
135
self .file_being_called_from : str = str (Path (sys ._getframe ().f_back .f_code .co_filename ).name ).replace ("." , "_" ) # noqa: SLF001
136
+ self .replay_test_file_path : Path | None = None
136
137
137
138
assert timeout is None or timeout > 0 , "Timeout should be greater than 0"
138
139
self .timeout = timeout
@@ -283,6 +284,7 @@ def __exit__(
283
284
284
285
with Path (test_file_path ).open ("w" , encoding = "utf8" ) as file :
285
286
file .write (replay_test )
287
+ self .replay_test_file_path = test_file_path
286
288
287
289
console .print (
288
290
f"Codeflash: Traced { self .trace_count } function calls successfully and replay test created at - { test_file_path } " ,
@@ -347,8 +349,7 @@ def tracer_logic(self, frame: FrameType, event: str) -> None: # noqa: PLR0911
347
349
try :
348
350
function_qualified_name = f"{ file_name } :{ code .co_qualname } "
349
351
except AttributeError :
350
- function_qualified_name = f"{ file_name } :{ (class_name + ':' if class_name else '' )} { code .co_name } "
351
-
352
+ function_qualified_name = f"{ file_name } :{ (class_name + '.' if class_name else '' )} { code .co_name } "
352
353
if function_qualified_name in self .ignored_qualified_functions :
353
354
return
354
355
if function_qualified_name not in self .function_count :
@@ -701,7 +702,7 @@ def print_stats(self, sort: str | int | tuple = -1) -> None:
701
702
border_style = "blue" ,
702
703
title = "[bold]Function Profile[/bold] (ordered by internal time)" ,
703
704
title_style = "cyan" ,
704
- caption = f"Showing top 25 of { len (self .stats )} functions" ,
705
+ caption = f"Showing top { min ( 25 , len ( self . stats )) } of { len (self .stats )} functions" ,
705
706
)
706
707
707
708
table .add_column ("Calls" , justify = "right" , style = "green" , width = 10 )
@@ -793,7 +794,7 @@ def runctx(self, cmd: str, global_vars: dict[str, Any], local_vars: dict[str, An
793
794
794
795
def main () -> ArgumentParser :
795
796
parser = ArgumentParser (allow_abbrev = False )
796
- parser .add_argument ("-o" , "--outfile" , dest = "outfile" , help = "Save trace to <outfile>" , required = True )
797
+ parser .add_argument ("-o" , "--outfile" , dest = "outfile" , help = "Save trace to <outfile>" , default = "codeflash.trace" )
797
798
parser .add_argument ("--only-functions" , help = "Trace only these functions" , nargs = "+" , default = None )
798
799
parser .add_argument (
799
800
"--max-function-count" ,
@@ -815,6 +816,7 @@ def main() -> ArgumentParser:
815
816
"with the codeflash config. Will be auto-discovered if not specified." ,
816
817
default = None ,
817
818
)
819
+ parser .add_argument ("--trace-only" , action = "store_true" , help = "Trace and create replay tests only, don't optimize" )
818
820
819
821
if not sys .argv [1 :]:
820
822
parser .print_usage ()
@@ -827,6 +829,7 @@ def main() -> ArgumentParser:
827
829
# to the output file at startup.
828
830
if args .outfile is not None :
829
831
args .outfile = Path (args .outfile ).resolve ()
832
+ outfile = args .outfile
830
833
831
834
if len (unknown_args ) > 0 :
832
835
if args .module :
@@ -848,14 +851,48 @@ def main() -> ArgumentParser:
848
851
"__cached__" : None ,
849
852
}
850
853
try :
851
- Tracer (
854
+ tracer = Tracer (
852
855
output = args .outfile ,
853
856
functions = args .only_functions ,
854
857
max_function_count = args .max_function_count ,
855
858
timeout = args .tracer_timeout ,
856
859
config_file_path = args .codeflash_config ,
857
860
command = " " .join (sys .argv ),
858
- ).runctx (code , globs , None )
861
+ )
862
+ tracer .runctx (code , globs , None )
863
+ replay_test_path = tracer .replay_test_file_path
864
+ if not args .trace_only and replay_test_path is not None :
865
+ del tracer
866
+
867
+ from codeflash .cli_cmds .cli import parse_args , process_pyproject_config
868
+ from codeflash .cli_cmds .cmd_init import CODEFLASH_LOGO
869
+ from codeflash .cli_cmds .console import paneled_text
870
+ from codeflash .telemetry import posthog_cf
871
+ from codeflash .telemetry .sentry import init_sentry
872
+
873
+ sys .argv = ["codeflash" , "--replay-test" , str (replay_test_path )]
874
+
875
+ args = parse_args ()
876
+ paneled_text (
877
+ CODEFLASH_LOGO ,
878
+ panel_args = {"title" : "https://codeflash.ai" , "expand" : False },
879
+ text_args = {"style" : "bold gold3" },
880
+ )
881
+
882
+ args = process_pyproject_config (args )
883
+ args .previous_checkpoint_functions = None
884
+ init_sentry (not args .disable_telemetry , exclude_errors = True )
885
+ posthog_cf .initialize_posthog (not args .disable_telemetry )
886
+
887
+ from codeflash .optimization import optimizer
888
+
889
+ optimizer .run_with_args (args )
890
+
891
+ # Delete the trace file and the replay test file if they exist
892
+ if outfile :
893
+ outfile .unlink (missing_ok = True )
894
+ if replay_test_path :
895
+ replay_test_path .unlink (missing_ok = True )
859
896
860
897
except BrokenPipeError as exc :
861
898
# Prevent "Exception ignored" during interpreter shutdown.
0 commit comments