safe changes

Author: KRRT7

codeflash/cli_cmds/cmd_init.py

@@ -141,12 +141,11 @@ def should_modify_pyproject_toml() -> bool:
     if "tests_root" not in config or config["tests_root"] is None or not Path(config["tests_root"]).is_dir():
         return True
 
-    create_toml = Confirm.ask(
+    return Confirm.ask(
         "✅ A valid Codeflash config already exists in this project. Do you want to re-configure it?",
         default=False,
         show_default=True,
     )
-    return create_toml
 
 
 def collect_setup_info() -> SetupInfo:

codeflash/code_utils/code_extractor.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import ast
-from pathlib import Path
 from typing import TYPE_CHECKING
 
 import libcst as cst
@@ -11,12 +10,15 @@
 from libcst.helpers import calculate_module_and_package
 
 from codeflash.cli_cmds.console import logger
-from codeflash.models.models import FunctionParent, FunctionSource
+from codeflash.models.models import FunctionParent
 
 if TYPE_CHECKING:
+    from pathlib import Path
+
     from libcst.helpers import ModuleNameAndPackage
 
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+    from codeflash.models.models import FunctionSource
 
 
 class FutureAliasedImportTransformer(cst.CSTTransformer):

codeflash/context/code_context_extractor.py

@@ -3,18 +3,15 @@
 import os
 from collections import defaultdict
 from itertools import chain
-from pathlib import Path
+from typing import TYPE_CHECKING, Optional
 
 import jedi
 import libcst as cst
 import tiktoken
-from jedi.api.classes import Name
-from libcst import CSTNode
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_extractor import add_needed_imports_from_module, find_preexisting_objects
 from codeflash.code_utils.code_utils import get_qualified_name, path_belongs_to_site_packages
-from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.models.models import (
     CodeContextType,
     CodeOptimizationContext,
@@ -24,6 +21,14 @@
 )
 from codeflash.optimization.function_context import belongs_to_function_qualified
 
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from jedi.api.classes import Name
+    from libcst import CSTNode
+
+    from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+
 
 def get_code_optimization_context(
     function_to_optimize: FunctionToOptimize,
@@ -75,7 +80,8 @@ def get_code_optimization_context(
     tokenizer = tiktoken.encoding_for_model("gpt-4o")
     final_read_writable_tokens = len(tokenizer.encode(final_read_writable_code))
     if final_read_writable_tokens > optim_token_limit:
-        raise ValueError("Read-writable code has exceeded token limit, cannot proceed")
+        msg = "Read-writable code has exceeded token limit, cannot proceed"
+        raise ValueError(msg)
 
     # Setup preexisting objects for code replacer
     preexisting_objects = set(
@@ -122,7 +128,8 @@ def get_code_optimization_context(
         testgen_context_code = testgen_code_markdown.code
         testgen_context_code_tokens = len(tokenizer.encode(testgen_context_code))
         if testgen_context_code_tokens > testgen_token_limit:
-            raise ValueError("Testgen code context has exceeded token limit, cannot proceed")
+            msg = "Testgen code context has exceeded token limit, cannot proceed"
+            raise ValueError(msg)
 
     return CodeOptimizationContext(
         testgen_context_code=testgen_context_code,
@@ -143,7 +150,7 @@ def extract_code_string_context_from_files(
     """Extract code context from files containing target functions and their helpers.
     This function processes two sets of files:
     1. Files containing the function to optimize (fto) and their first-degree helpers
-    2. Files containing only helpers of helpers (with no overlap with the first set)
+    2. Files containing only helpers of helpers (with no overlap with the first set).
 
     For each file, it extracts relevant code based on the specified context type, adds necessary
     imports, and combines them.
@@ -358,18 +365,18 @@ def get_function_to_optimize_as_function_source(
             and name.name == function_to_optimize.function_name
             and get_qualified_name(name.module_name, name.full_name) == function_to_optimize.qualified_name
         ):
-            function_source = FunctionSource(
+            return FunctionSource(
                 file_path=function_to_optimize.file_path,
                 qualified_name=function_to_optimize.qualified_name,
                 fully_qualified_name=name.full_name,
                 only_function_name=name.name,
                 source_code=name.get_line_code(),
                 jedi_definition=name,
             )
-            return function_source
 
+    msg = f"Could not find function {function_to_optimize.function_name} in {function_to_optimize.file_path}"
     raise ValueError(
-        f"Could not find function {function_to_optimize.function_name} in {function_to_optimize.file_path}"
+        msg
     )
 
 
@@ -436,7 +443,7 @@ def get_section_names(node: cst.CSTNode) -> list[str]:
 
 
 def remove_docstring_from_body(indented_block: cst.IndentedBlock) -> cst.CSTNode:
-    """Removes the docstring from an indented block if it exists"""
+    """Removes the docstring from an indented block if it exists."""
     if not isinstance(indented_block.body[0], cst.SimpleStatementLine):
         return indented_block
     first_stmt = indented_block.body[0].body[0]
@@ -449,10 +456,12 @@ def parse_code_and_prune_cst(
     code: str,
     code_context_type: CodeContextType,
     target_functions: set[str],
-    helpers_of_helper_functions: set[str] = set(),
+    helpers_of_helper_functions: Optional[set[str]] = None,
     remove_docstrings: bool = False,
 ) -> str:
     """Create a read-only version of the code by parsing and filtering the code to keep only class contextual information, and other module scoped variables."""
+    if helpers_of_helper_functions is None:
+        helpers_of_helper_functions = set()
     module = cst.parse_module(code)
     if code_context_type == CodeContextType.READ_WRITABLE:
         filtered_node, found_target = prune_cst_for_read_writable_code(module, target_functions)
@@ -465,10 +474,12 @@ def parse_code_and_prune_cst(
             module, target_functions, helpers_of_helper_functions, remove_docstrings=remove_docstrings
         )
     else:
-        raise ValueError(f"Unknown code_context_type: {code_context_type}")
+        msg = f"Unknown code_context_type: {code_context_type}"
+        raise ValueError(msg)
 
     if not found_target:
-        raise ValueError("No target functions found in the provided code")
+        msg = "No target functions found in the provided code"
+        raise ValueError(msg)
     if filtered_node and isinstance(filtered_node, cst.Module):
         return str(filtered_node.code)
     return ""
@@ -500,7 +511,8 @@ def prune_cst_for_read_writable_code(
             return None, False
         # Assuming always an IndentedBlock
         if not isinstance(node.body, cst.IndentedBlock):
-            raise ValueError("ClassDef body is not an IndentedBlock")
+            msg = "ClassDef body is not an IndentedBlock"
+            raise ValueError(msg)
         class_prefix = f"{prefix}.{node.name.value}" if prefix else node.name.value
         new_body = []
         found_target = False
@@ -593,7 +605,8 @@ def prune_cst_for_read_only_code(
             return None, False
         # Assuming always an IndentedBlock
         if not isinstance(node.body, cst.IndentedBlock):
-            raise ValueError("ClassDef body is not an IndentedBlock")
+            msg = "ClassDef body is not an IndentedBlock"
+            raise ValueError(msg)
 
         class_prefix = f"{prefix}.{node.name.value}" if prefix else node.name.value
 
@@ -698,7 +711,8 @@ def prune_cst_for_testgen_code(
             return None, False
         # Assuming always an IndentedBlock
         if not isinstance(node.body, cst.IndentedBlock):
-            raise ValueError("ClassDef body is not an IndentedBlock")
+            msg = "ClassDef body is not an IndentedBlock"
+            raise ValueError(msg)
 
         class_prefix = f"{prefix}.{node.name.value}" if prefix else node.name.value
 

codeflash/discovery/pytest_new_process_discovery.py

@@ -16,7 +16,7 @@ def pytest_collection_finish(self, session) -> None:
         collected_tests.extend(session.items)
         pytest_rootdir = session.config.rootdir
 
-    def pytest_collection_modifyitems(config, items):
+    def pytest_collection_modifyitems(self, items) -> None:
         skip_benchmark = pytest.mark.skip(reason="Skipping benchmark tests")
         for item in items:
             if "benchmark" in item.fixturenames:

codeflash/models/models.py

@@ -11,7 +11,7 @@
 import enum
 import re
 import sys
-from collections.abc import Collection, Iterator
+from collections.abc import Collection
 from enum import Enum, IntEnum
 from pathlib import Path
 from re import Pattern

codeflash/optimization/function_context.py

@@ -1,9 +1,12 @@
 from __future__ import annotations
 
-from jedi.api.classes import Name
+from typing import TYPE_CHECKING
 
 from codeflash.code_utils.code_utils import get_qualified_name
 
+if TYPE_CHECKING:
+    from jedi.api.classes import Name
+
 
 def belongs_to_method(name: Name, class_name: str, method_name: str) -> bool:
     """Check if the given name belongs to the specified method."""
@@ -14,9 +17,8 @@ def belongs_to_function(name: Name, function_name: str) -> bool:
     """Check if the given jedi Name is a direct child of the specified function."""
     if name.name == function_name:  # Handles function definition and recursive function calls
         return False
-    if name := name.parent():
-        if name.type == "function":
-            return name.name == function_name
+    if (name := name.parent()) and name.type == "function":
+        return name.name == function_name
     return False
 
 
@@ -34,9 +36,8 @@ def belongs_to_function_qualified(name: Name, qualified_function_name: str) -> b
         if get_qualified_name(name.module_name, name.full_name) == qualified_function_name:
             # Handles function definition and recursive function calls
             return False
-        if name := name.parent():
-            if name.type == "function":
-                return get_qualified_name(name.module_name, name.full_name) == qualified_function_name
+        if (name := name.parent()) and name.type == "function":
+            return get_qualified_name(name.module_name, name.full_name) == qualified_function_name
         return False
     except ValueError:
         return False

codeflash/optimization/function_optimizer.py

@@ -41,7 +41,6 @@
 from codeflash.code_utils.static_analysis import get_first_top_level_function_or_method_ast
 from codeflash.code_utils.time_utils import humanize_runtime
 from codeflash.context import code_context_extractor
-from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.either import Failure, Success, is_successful
 from codeflash.models.ExperimentMetadata import ExperimentMetadata
 from codeflash.models.models import (
@@ -74,6 +73,7 @@
 if TYPE_CHECKING:
     from argparse import Namespace
 
+    from codeflash.discovery.functions_to_optimize import FunctionToOptimize
     from codeflash.either import Result
     from codeflash.models.models import CoverageData, FunctionSource, OptimizedCandidate
     from codeflash.verification.verification_utils import TestConfig

codeflash/tracing/replay_test.py

@@ -2,11 +2,15 @@
 
 import sqlite3
 import textwrap
-from collections.abc import Generator
-from typing import Any, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
-from codeflash.discovery.functions_to_optimize import FunctionProperties, inspect_top_level_functions_or_methods
-from codeflash.tracing.tracing_utils import FunctionModules
+from codeflash.discovery.functions_to_optimize import inspect_top_level_functions_or_methods
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+    from codeflash.discovery.functions_to_optimize import FunctionProperties
+    from codeflash.tracing.tracing_utils import FunctionModules
 
 
 def get_next_arg_and_return(

codeflash/verification/concolic_testing.py

@@ -3,19 +3,23 @@
 import ast
 import subprocess
 import tempfile
-from argparse import Namespace
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 from codeflash.cli_cmds.console import console, logger
 from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE
 from codeflash.code_utils.concolic_utils import clean_concolic_tests
 from codeflash.code_utils.static_analysis import has_typed_parameters
 from codeflash.discovery.discover_unit_tests import discover_unit_tests
-from codeflash.discovery.functions_to_optimize import FunctionToOptimize
-from codeflash.models.models import FunctionCalledInTest
 from codeflash.telemetry.posthog_cf import ph
 from codeflash.verification.verification_utils import TestConfig
 
+if TYPE_CHECKING:
+    from argparse import Namespace
+
+    from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+    from codeflash.models.models import FunctionCalledInTest
+
 
 def generate_concolic_tests(
     test_cfg: TestConfig, args: Namespace, function_to_optimize: FunctionToOptimize, function_to_optimize_ast: ast.AST

codeflash/verification/instrument_codeflash_capture.py

@@ -2,11 +2,14 @@
 
 import ast
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import isort
 
 from codeflash.code_utils.code_utils import get_run_tmp_file
-from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+
+if TYPE_CHECKING:
+    from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 
 
 def instrument_codeflash_capture(
@@ -109,7 +112,7 @@ def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
             func=ast.Name(id="codeflash_capture", ctx=ast.Load()),
             args=[],
             keywords=[
-                ast.keyword(arg="function_name", value=ast.Constant(value=".".join([node.name, "__init__"]))),
+                ast.keyword(arg="function_name", value=ast.Constant(value=f"{node.name}.__init__")),
                 ast.keyword(arg="tmp_dir_path", value=ast.Constant(value=self.tmp_dir_path)),
                 ast.keyword(arg="tests_root", value=ast.Constant(value=str(self.tests_root))),
                 ast.keyword(arg="is_fto", value=ast.Constant(value=self.is_fto)),