organized imports

codeflash/models/models.py

@@ -1,15 +1,14 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Optional, cast
+from typing import TYPE_CHECKING
 
 from rich.tree import Tree
 
-from codeflash.cli_cmds.console import DEBUG_MODE, logger
+from codeflash.cli_cmds.console import DEBUG_MODE
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
 import enum
-import json
 import re
 import sys
 from collections.abc import Collection, Iterator
@@ -18,19 +17,12 @@ from pathlib import Path
 from re import Pattern
 from typing import Annotated, Any, Optional, Union, cast
 
-import sentry_sdk
-from coverage.exceptions import NoDataError
 from jedi.api.classes import Name
 from pydantic import AfterValidator, BaseModel, ConfigDict, Field
 from pydantic.dataclasses import dataclass
 
 from codeflash.cli_cmds.console import console, logger
 from codeflash.code_utils.code_utils import validate_python_code
-from codeflash.code_utils.coverage_utils import (
-    build_fully_qualified_name,
-    extract_dependent_function,
-    generate_candidates,
-)
 from codeflash.code_utils.env_utils import is_end_to_end
 from codeflash.verification.comparator import comparator
 

codeflash/optimization/function_optimizer.py

@@ -26,8 +26,8 @@ from codeflash.code_utils.code_utils import (
     cleanup_paths,
     file_name_from_test_module_name,
     get_run_tmp_file,
-    module_name_from_file_path,
     has_any_async_functions,
+    module_name_from_file_path,
 )
 from codeflash.code_utils.config_consts import (
     INDIVIDUAL_TESTCASE_TIMEOUT,
@@ -57,7 +57,7 @@ from codeflash.models.models import (
     TestFiles,
     TestingMode,
     TestResults,
-    TestType
+    TestType,
 )
 from codeflash.result.create_pr import check_create_pr, existing_tests_source_for
 from codeflash.result.critic import coverage_critic, performance_gain, quantity_of_tests_critic, speedup_critic

codeflash/optimization/optimizer.py

@@ -16,7 +16,7 @@ from codeflash.code_utils.static_analysis import analyze_imported_modules, get_f
 from codeflash.discovery.discover_unit_tests import discover_unit_tests
 from codeflash.discovery.functions_to_optimize import get_functions_to_optimize
 from codeflash.either import is_successful
-from codeflash.models.models import ValidCode, TestType
+from codeflash.models.models import TestType, ValidCode
 from codeflash.optimization.function_optimizer import FunctionOptimizer
 from codeflash.telemetry.posthog_cf import ph
 from codeflash.verification.verification_utils import TestConfig

codeflash/result/critic.py

@@ -1,9 +1,15 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils import env_utils
 from codeflash.code_utils.config_consts import COVERAGE_THRESHOLD, MIN_IMPROVEMENT_THRESHOLD
-from codeflash.models.models import CoverageData, OptimizedCandidateResult, TestType
+from codeflash.models.models import TestType
+
+if TYPE_CHECKING:
+    from codeflash.models.models import CoverageData, OptimizedCandidateResult
+
 
 def performance_gain(*, original_runtime_ns: int, optimized_runtime_ns: int) -> float:
     """Calculate the performance gain of an optimized code over the original code.

codeflash/verification/coverage_utils.py

@@ -1,23 +1,24 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Optional, cast
-
-
 import json
-from collections.abc import Collection, Iterator
-from pathlib import Path
-from typing import Annotated, Any, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Union
 
 import sentry_sdk
 from coverage.exceptions import NoDataError
 
-from codeflash.cli_cmds.console import console, logger
+from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.coverage_utils import (
     build_fully_qualified_name,
     extract_dependent_function,
     generate_candidates,
 )
-from codeflash.models.models import CoverageData, CodeOptimizationContext, CoverageStatus, FunctionCoverage
+from codeflash.models.models import CoverageData, CoverageStatus, FunctionCoverage
+
+if TYPE_CHECKING:
+    from collections.abc import Collection
+    from pathlib import Path
+
+    from codeflash.models.models import CodeOptimizationContext
 
 
 class CoverageUtils:

codeflash/verification/equivalence.py

@@ -1,9 +1,8 @@
-import difflib
 import sys
 
-from codeflash.cli_cmds.console import console, logger
-from codeflash.verification.comparator import comparator
+from codeflash.cli_cmds.console import logger
 from codeflash.models.models import TestResults, TestType, VerificationType
+from codeflash.verification.comparator import comparator
 
 INCREASED_RECURSION_LIMIT = 5000
 

codeflash/verification/parse_test_output.py

@@ -20,21 +20,13 @@ from codeflash.code_utils.code_utils import (
     module_name_from_file_path,
 )
 from codeflash.discovery.discover_unit_tests import discover_parameters_unittest
-from codeflash.models.models import (
-    CoverageData,
-    TestFiles,
-    FunctionTestInvocation,
-    InvocationId,
-    TestResults,
-    TestType,
-    VerificationType,
-)
+from codeflash.models.models import FunctionTestInvocation, InvocationId, TestResults, TestType, VerificationType
 from codeflash.verification.coverage_utils import CoverageUtils
 
 if TYPE_CHECKING:
     import subprocess
 
-    from codeflash.models.models import CodeOptimizationContext
+    from codeflash.models.models import CodeOptimizationContext, CoverageData, TestFiles
     from codeflash.verification.verification_utils import TestConfig
 
 

tests/test_comparator.py

@@ -12,9 +12,9 @@ import pydantic
 import pytest
 
 from codeflash.either import Failure, Success
+from codeflash.models.models import FunctionTestInvocation, InvocationId, TestResults, TestType
 from codeflash.verification.comparator import comparator
 from codeflash.verification.equivalence import compare_test_results
-from codeflash.models.models import FunctionTestInvocation, InvocationId, TestResults, TestType
 
 
 def test_basic_python_objects() -> None:

tests/test_critic.py

@@ -8,10 +8,13 @@ from codeflash.models.models import (
     CoverageData,
     CoverageStatus,
     FunctionCoverage,
+    FunctionTestInvocation,
+    InvocationId,
     OptimizedCandidateResult,
+    TestResults,
+    TestType,
 )
 from codeflash.result.critic import coverage_critic, performance_gain, quantity_of_tests_critic, speedup_critic
-from codeflash.models.models import FunctionTestInvocation, InvocationId, TestResults, TestType
 
 
 def test_performance_gain() -> None:

tests/test_instrument_tests.py

@@ -13,7 +13,15 @@ from codeflash.code_utils.instrument_existing_tests import (
     inject_profiling_into_existing_test,
 )
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
-from codeflash.models.models import CodePosition, FunctionParent, TestFile, TestFiles, TestingMode, TestsInFile, TestType
+from codeflash.models.models import (
+    CodePosition,
+    FunctionParent,
+    TestFile,
+    TestFiles,
+    TestingMode,
+    TestsInFile,
+    TestType,
+)
 from codeflash.optimization.function_optimizer import FunctionOptimizer
 from codeflash.verification.verification_utils import TestConfig
 

tests/test_merge_test_results.py

@@ -1,5 +1,5 @@
-from codeflash.verification.parse_test_output import merge_test_results
 from codeflash.models.models import FunctionTestInvocation, InvocationId, TestResults, TestType
+from codeflash.verification.parse_test_output import merge_test_results
 
 
 def test_merge_test_results_1():