fix formatting issues

Author: HeshamHM28

codeflash/api/cfapi.py

@@ -14,14 +14,15 @@
 
 from codeflash.cli_cmds.console import console, logger
 from codeflash.code_utils.env_utils import ensure_codeflash_api_key, get_codeflash_api_key, get_pr_number
-from codeflash.code_utils.git_utils import get_repo_owner_and_name, git_root_dir, get_current_branch
-from codeflash.result.explanation import Explanation
-from codeflash.version import __version__
+from codeflash.code_utils.git_utils import get_current_branch, get_repo_owner_and_name, git_root_dir
 from codeflash.github.PrComment import FileDiffContent, PrComment
+from codeflash.version import __version__
 
 if TYPE_CHECKING:
     from requests import Response
 
+    from codeflash.result.explanation import Explanation
+
 from packaging import version
 
 if os.environ.get("CODEFLASH_CFAPI_SERVER", default="prod").lower() == "local":
@@ -175,8 +176,9 @@ def create_pr(
     }
     return make_cfapi_request(endpoint="/create-pr", method="POST", payload=payload)
 
+
 def create_staging(
-    original_code:str,
+    original_code: str,
     new_code: str,
     explanation: Explanation,
     existing_tests_source: str,
@@ -204,27 +206,28 @@ def create_staging(
         for p in original_code
     }
     payload = {
-        "baseBranch":  get_current_branch(),
+        "baseBranch": get_current_branch(),
         "diffContents": build_file_changes,
-        "prCommentFields":   PrComment(
-        optimization_explanation=explanation.explanation_message(),
-        best_runtime=explanation.best_runtime_ns,
-        original_runtime=explanation.original_runtime_ns,
-        function_name=explanation.function_name,
-        relative_file_path=relative_path,
-        speedup_x=explanation.speedup_x,
-        speedup_pct=explanation.speedup_pct,
-        winning_behavioral_test_results=explanation.winning_behavioral_test_results,
-        winning_benchmarking_test_results=explanation.winning_benchmarking_test_results,
-        benchmark_details=explanation.benchmark_details,
-    ).to_json(),
+        "prCommentFields": PrComment(
+            optimization_explanation=explanation.explanation_message(),
+            best_runtime=explanation.best_runtime_ns,
+            original_runtime=explanation.original_runtime_ns,
+            function_name=explanation.function_name,
+            relative_file_path=relative_path,
+            speedup_x=explanation.speedup_x,
+            speedup_pct=explanation.speedup_pct,
+            winning_behavioral_test_results=explanation.winning_behavioral_test_results,
+            winning_benchmarking_test_results=explanation.winning_benchmarking_test_results,
+            benchmark_details=explanation.benchmark_details,
+        ).to_json(),
         "existingTests": existing_tests_source,
         "generatedTests": generated_original_test_source,
         "traceId": function_trace_id,
         "coverage_message": coverage_message,
     }
     return make_cfapi_request(endpoint="/create-staging", method="POST", payload=payload)
 
+
 def is_github_app_installed_on_repo(owner: str, repo: str) -> bool:
     """Check if the Codeflash GitHub App is installed on the specified repository.
 

codeflash/cli_cmds/cli.py

@@ -47,11 +47,7 @@ def parse_args() -> Namespace:
     parser.add_argument(
         "--no-pr", action="store_true", help="Do not create a PR for the optimization, only update the code locally."
     )
-    parser.add_argument(
-        "--staging-review",
-        action="store_true",
-        help="Upload optimizations to staging for review"
-    )
+    parser.add_argument("--staging-review", action="store_true", help="Upload optimizations to staging for review")
     parser.add_argument(
         "--verify-setup",
         action="store_true",

codeflash/optimization/function_optimizer.py

@@ -19,7 +19,7 @@
 from rich.tree import Tree
 
 from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient
-from codeflash.api.cfapi import add_code_context_hash, mark_optimization_success, create_staging
+from codeflash.api.cfapi import add_code_context_hash, create_staging, mark_optimization_success
 from codeflash.benchmarking.utils import process_benchmark_data
 from codeflash.cli_cmds.console import code_print, console, logger, progress_bar
 from codeflash.code_utils import env_utils
@@ -997,28 +997,50 @@ def find_and_process_best_optimization(
                 original_code_combined[explanation.file_path] = self.function_to_optimize_source_code
                 new_code_combined = new_helper_code.copy()
                 new_code_combined[explanation.file_path] = new_code
-                self.process_review(original_code_baseline, best_optimization, generated_tests,
-                              test_functions_to_remove, concolic_test_str, original_code_combined, new_code_combined,
-                              explanation, function_to_all_tests, exp_type, original_helper_code)
+                self.process_review(
+                    original_code_baseline,
+                    best_optimization,
+                    generated_tests,
+                    test_functions_to_remove,
+                    concolic_test_str,
+                    original_code_combined,
+                    new_code_combined,
+                    explanation,
+                    function_to_all_tests,
+                    exp_type,
+                    original_helper_code,
+                )
                 self.log_successful_optimization(explanation, generated_tests, exp_type)
         return best_optimization
 
     def process_review(
-        self, original_code_baseline, best_optimization, generated_tests,
-        test_functions_to_remove, concolic_test_str, original_code_combined,
-        new_code_combined, explanation, function_to_all_tests, exp_type, original_helper_code
-    ):
+        self,
+        original_code_baseline: OriginalCodeBaseline,
+        best_optimization: BestOptimization,
+        generated_tests: GeneratedTestsList,
+        test_functions_to_remove: list[str],
+        concolic_test_str: str | None,
+        original_code_combined: dict[Path, str],
+        new_code_combined: dict[Path, str],
+        explanation: Explanation,
+        function_to_all_tests: dict[str, set[FunctionCalledInTest]],
+        exp_type: str,
+        original_helper_code: dict[Path, str],
+    ) -> None:
         coverage_message = (
             original_code_baseline.coverage_results.build_message()
-            if original_code_baseline.coverage_results else "Coverage data not available"
+            if original_code_baseline.coverage_results
+            else "Coverage data not available"
         )
 
         generated_tests = remove_functions_from_generated_tests(
             generated_tests=generated_tests, test_functions_to_remove=test_functions_to_remove
         )
 
         original_runtime_by_test = original_code_baseline.benchmarking_test_results.usable_runtime_data_by_test_case()
-        optimized_runtime_by_test = best_optimization.winning_benchmarking_test_results.usable_runtime_data_by_test_case()
+        optimized_runtime_by_test = (
+            best_optimization.winning_benchmarking_test_results.usable_runtime_data_by_test_case()
+        )
 
         generated_tests = add_runtime_comments_to_generated_tests(
             self.test_cfg, generated_tests, original_runtime_by_test, optimized_runtime_by_test
@@ -1038,15 +1060,17 @@ def process_review(
             optimized_runtimes_all=optimized_runtime_by_test,
         )
 
-        data = dict(
-            original_code=original_code_combined,
-            new_code=new_code_combined,
-            explanation=explanation,
-            existing_tests_source=existing_tests,
-            generated_original_test_source=generated_tests_str,
-            function_trace_id=self.function_trace_id[:-4] + exp_type if self.experiment_id else self.function_trace_id,
-            coverage_message=coverage_message,
-        )
+        data = {
+            "original_code": original_code_combined,
+            "new_code": new_code_combined,
+            "explanation": explanation,
+            "existing_tests_source": existing_tests,
+            "generated_original_test_source": generated_tests_str,
+            "function_trace_id": self.function_trace_id[:-4] + exp_type
+            if self.experiment_id
+            else self.function_trace_id,
+            "coverage_message": coverage_message,
+        }
 
         if not self.args.no_pr and not self.args.staging_review:
             data["git_remote"] = self.args.git_remote
@@ -1059,13 +1083,12 @@ def process_review(
                 trace_id=self.function_trace_id, is_optimization_found=best_optimization is not None
             )
 
-        if (not self.args.no_pr) or not self.args.staging_review:
-            if self.args.all or env_utils.get_pr_number() or (self.args.file and not self.args.function):
-                self.write_code_and_helpers(
-                    self.function_to_optimize_source_code,
-                    original_helper_code,
-                    self.function_to_optimize.file_path,
-                )
+        if ((not self.args.no_pr) or not self.args.staging_review) and (
+            self.args.all or env_utils.get_pr_number() or (self.args.file and not self.args.function)
+        ):
+            self.write_code_and_helpers(
+                self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path
+            )
 
     def establish_original_code_baseline(
         self,