perf: defer libcst, Rich, comparator imports in models.py

Move libcst, rich.tree.Tree, console, comparator, code_utils, registry,
lsp.helpers, and LspMarkdownMessage from module-level to the methods that
use them. Only pydantic and TestType remain at module level (needed for
class definitions).

models.py import: 633ms → 125ms on Azure Standard_D4s_v5.

Author: KRRT7

codeflash/models/models.py

@@ -1,37 +1,23 @@
 from __future__ import annotations
 
-from collections import Counter, defaultdict
-from functools import lru_cache
-from typing import TYPE_CHECKING
-
-import libcst as cst
-from rich.tree import Tree
-
-from codeflash.cli_cmds.console import DEBUG_MODE, lsp_log
-from codeflash.languages.registry import get_language_support
-from codeflash.lsp.helpers import is_LSP_enabled, report_to_markdown_table
-from codeflash.lsp.lsp_message import LspMarkdownMessage
-from codeflash.models.test_type import TestType
-
-if TYPE_CHECKING:
-    from collections.abc import Iterator
-
 import enum
 import re
 import sys
+from collections import Counter, defaultdict
 from collections.abc import Collection
 from enum import Enum, IntEnum
+from functools import lru_cache
 from pathlib import Path
 from re import Pattern
-from typing import Any, NamedTuple, Optional, cast
+from typing import TYPE_CHECKING, Any, NamedTuple, Optional, cast
 
 from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, ValidationError, model_validator
 from pydantic.dataclasses import dataclass
 
-from codeflash.cli_cmds.console import console, logger
-from codeflash.code_utils.code_utils import diff_length, module_name_from_file_path, validate_python_code
-from codeflash.code_utils.env_utils import is_end_to_end
-from codeflash.verification.comparator import comparator
+from codeflash.models.test_type import TestType
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
 
 
 @dataclass(frozen=True)
@@ -254,6 +240,8 @@ class CodeString(BaseModel):
     def validate_code_syntax(self) -> CodeString:
         """Validate code syntax for the specified language."""
         if self.language == "python":
+            from codeflash.code_utils.code_utils import validate_python_code  # noqa: PLC0415
+
             validate_python_code(self.code)
         else:
             from codeflash.languages.registry import get_language_support
@@ -267,6 +255,8 @@ def validate_code_syntax(self) -> CodeString:
 
 def get_comment_prefix(file_path: Path) -> str:
     """Get the comment prefix for a given language."""
+    from codeflash.languages.registry import get_language_support  # noqa: PLC0415
+
     support = get_language_support(file_path)
     return support.comment_prefix
 
@@ -565,6 +555,8 @@ def handle_duplicate_candidate(
         self.optimizations_post[past_opt_id] = self.ast_code_to_id[normalized_code]["shorter_source_code"].markdown
 
         # Update to shorter code if this candidate has a shorter diff
+        from codeflash.code_utils.code_utils import diff_length  # noqa: PLC0415
+
         new_diff_len = diff_length(candidate.source_code.flat, original_flat_code)
         if new_diff_len < self.ast_code_to_id[normalized_code]["diff_len"]:
             self.ast_code_to_id[normalized_code]["shorter_source_code"] = candidate.source_code
@@ -574,6 +566,8 @@ def register_new_candidate(
         self, normalized_code: str, candidate: OptimizedCandidate, original_flat_code: str
     ) -> None:
         """Register a new candidate that hasn't been seen before."""
+        from codeflash.code_utils.code_utils import diff_length  # noqa: PLC0415
+
         self.ast_code_to_id[normalized_code] = {
             "optimization_id": candidate.optimization_id,
             "shorter_source_code": candidate.source_code,
@@ -668,7 +662,10 @@ def build_message(self) -> str:
         return f"{self.coverage:.1f}%"
 
     def log_coverage(self) -> None:
-        from rich.tree import Tree
+        from rich.tree import Tree  # noqa: PLC0415
+
+        from codeflash.cli_cmds.console import console, logger  # noqa: PLC0415
+        from codeflash.code_utils.env_utils import is_end_to_end  # noqa: PLC0415
 
         tree = Tree("Test Coverage Results")
         tree.add(f"Main Function: {self.main_func_coverage.name}: {self.coverage:.2f}%")
@@ -769,12 +766,16 @@ def test_fn_qualified_name(self) -> str:
         )
 
     def find_func_in_class(self, class_node: cst.ClassDef, func_name: str) -> Optional[cst.FunctionDef]:
+        import libcst as cst  # noqa: PLC0415
+
         for stmt in class_node.body.body:
             if isinstance(stmt, cst.FunctionDef) and stmt.name.value == func_name:
                 return stmt
         return None
 
     def get_src_code(self, test_path: Path) -> Optional[str]:
+        import libcst as cst  # noqa: PLC0415
+
         if not test_path.exists():
             return None
         try:
@@ -856,6 +857,8 @@ def add(self, function_test_invocation: FunctionTestInvocation) -> None:
         unique_id = function_test_invocation.unique_invocation_loop_id
         test_result_idx = self.test_result_idx
         if unique_id in test_result_idx:
+            from codeflash.cli_cmds.console import DEBUG_MODE, logger  # noqa: PLC0415
+
             if DEBUG_MODE:
                 logger.warning(f"Test result with id {unique_id} already exists. SKIPPING")
             return
@@ -876,6 +879,8 @@ def group_by_benchmarks(
         self, benchmark_keys: list[BenchmarkKey], benchmark_replay_test_dir: Path, project_root: Path
     ) -> dict[BenchmarkKey, TestResults]:
         """Group TestResults by benchmark for calculating improvements for each benchmark."""
+        from codeflash.code_utils.code_utils import module_name_from_file_path  # noqa: PLC0415
+
         test_results_by_benchmark = defaultdict(TestResults)
         benchmark_module_path = {}
         for benchmark_key in benchmark_keys:
@@ -929,9 +934,17 @@ def report_to_string(report: dict[TestType, dict[str, int]]) -> str:
 
     @staticmethod
     def report_to_tree(report: dict[TestType, dict[str, int]], title: str) -> Tree:
+        from rich.tree import Tree  # noqa: PLC0415
+
+        from codeflash.lsp.helpers import is_LSP_enabled  # noqa: PLC0415
+
         tree = Tree(title)
 
         if is_LSP_enabled():
+            from codeflash.cli_cmds.console import lsp_log  # noqa: PLC0415
+            from codeflash.lsp.helpers import report_to_markdown_table  # noqa: PLC0415
+            from codeflash.lsp.lsp_message import LspMarkdownMessage  # noqa: PLC0415
+
             # Build markdown table
             markdown = report_to_markdown_table(report, title)
             lsp_log(LspMarkdownMessage(markdown=markdown))
@@ -946,6 +959,8 @@ def report_to_tree(report: dict[TestType, dict[str, int]], title: str) -> Tree:
         return tree
 
     def usable_runtime_data_by_test_case(self) -> dict[InvocationId, list[int]]:
+        from codeflash.cli_cmds.console import logger  # noqa: PLC0415
+
         # Efficient single traversal, directly accumulating into a dict.
         # can track mins here and only sums can be return in total_passed_runtime
         by_id: dict[InvocationId, list[int]] = {}
@@ -1025,6 +1040,8 @@ def __bool__(self) -> bool:
         return bool(self.test_results)
 
     def __eq__(self, other: object) -> bool:
+        from codeflash.verification.comparator import comparator  # noqa: PLC0415
+
         # Unordered comparison
         if type(self) is not type(other):
             return False