Fixing more bugs, testing live...

Also updated CODEFLASH_DISABLE_TELEMETRY behavior

Author: zomglings

codeflash/main.py

@@ -22,8 +22,10 @@ def main() -> None:
     )
     args = parse_args()
 
+    disable_telemetry_env = os.environ.get("CODEFLASH_DISABLE_TELEMETRY", "").lower() in {"true", "t", "1", "yes", "y"}
+
     if args.command:
-        disable_telemetry = os.environ.get("CODEFLASH_DISABLE_TELEMETRY", "").lower() in {"true", "t", "1", "yes", "y"}
+        disable_telemetry = disable_telemetry_env
         if (not disable_telemetry) and args.config_file and Path.exists(args.config_file):
             pyproject_config, _ = parse_config_file(args.config_file)
             disable_telemetry = pyproject_config.get("disable_telemetry", False)
@@ -32,14 +34,16 @@ def main() -> None:
         args.func()
     elif args.verify_setup:
         args = process_pyproject_config(args)
-        init_sentry(not args.disable_telemetry, exclude_errors=True)
-        posthog_cf.initialize_posthog(not args.disable_telemetry)
+        disable_telemetry = args.disable_telemetry or disable_telemetry_env
+        init_sentry(not disable_telemetry, exclude_errors=True)
+        posthog_cf.initialize_posthog(not disable_telemetry)
         ask_run_end_to_end_test(args)
     else:
         args = process_pyproject_config(args)
         args.previous_checkpoint_functions = ask_should_use_checkpoint_get_functions(args)
-        init_sentry(not args.disable_telemetry, exclude_errors=True)
-        posthog_cf.initialize_posthog(not args.disable_telemetry)
+        disable_telemetry = args.disable_telemetry or disable_telemetry_env
+        init_sentry(not disable_telemetry, exclude_errors=True)
+        posthog_cf.initialize_posthog(not disable_telemetry)
         optimizer.run_with_args(args)
 
 

codeflash/optimization/function_optimizer.py

@@ -609,12 +609,14 @@ def reformat_code_and_helpers(
             # It is important to sort in descending order so that the index arithmetic remains simple as we modify new_code
             code_ranges_unformatted.sort(key=lambda range: range[0], reverse=True)
             code_ranges_formatted.sort(key=lambda range: range[0], reverse=True)
-            new_code = unformatted_code
+            formatted_code_lines = formatted_code.split("\n")
+            new_code_lines = unformatted_code.split("\n")
             for range_0, range_1 in zip(code_ranges_unformatted, code_ranges_formatted):
                 range_0_0, range_0_1 = range_0
                 range_1_0, range_1_1 = range_1
-                new_code = new_code[:range_0_0] + new_code[range_1_0:range_1_1 + 1] + new_code[range_0_1 + 1]
-
+                new_code_lines = new_code_lines[:range_0_0] + formatted_code_lines[range_1_0:range_1_1 + 1] + new_code_lines[range_0_1 + 1:]
+            new_code = "\n".join(new_code_lines)
+            breakpoint()
             path.write_text(new_code, encoding="utf8")
 
             if should_sort_imports: