cleanup fixes

Author: KRRT7

codeflash/api/aiservice.py

@@ -73,7 +73,7 @@ def make_ai_service_request(
         # response.raise_for_status()  # Will raise an HTTPError if the HTTP request returned an unsuccessful status code
         return response
 
-    def optimize_python_code(
+    def optimize_python_code(  # noqa: D417
         self,
         source_code: str,
         dependency_code: str,
@@ -139,7 +139,7 @@ def optimize_python_code(
         console.rule()
         return []
 
-    def optimize_python_code_line_profiler(
+    def optimize_python_code_line_profiler(  # noqa: D417
         self,
         source_code: str,
         dependency_code: str,
@@ -208,7 +208,7 @@ def optimize_python_code_line_profiler(
         console.rule()
         return []
 
-    def log_results(
+    def log_results(  # noqa: D417
         self,
         function_trace_id: str,
         speedup_ratio: dict[str, float | None] | None,
@@ -240,7 +240,7 @@ def log_results(
         except requests.exceptions.RequestException as e:
             logger.exception(f"Error logging features: {e}")
 
-    def generate_regression_tests(
+    def generate_regression_tests(  # noqa: D417
         self,
         source_code_being_tested: str,
         function_to_optimize: FunctionToOptimize,
@@ -307,7 +307,7 @@ def generate_regression_tests(
             error = response.json()["error"]
             logger.error(f"Error generating tests: {response.status_code} - {error}")
             ph("cli-testgen-error-response", {"response_status_code": response.status_code, "error": error})
-            return None
+            return None  # noqa: TRY300
         except Exception:
             logger.error(f"Error generating tests: {response.status_code} - {response.text}")
             ph("cli-testgen-error-response", {"response_status_code": response.status_code, "error": response.text})

codeflash/benchmarking/codeflash_trace.py

@@ -4,7 +4,7 @@
 import sqlite3
 import threading
 import time
-from typing import Callable
+from typing import Any, Callable
 
 from codeflash.picklepatch.pickle_patcher import PicklePatcher
 
@@ -103,7 +103,7 @@ def __call__(self, func: Callable) -> Callable:
         func_id = (func.__module__, func.__name__)
 
         @functools.wraps(func)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args, **kwargs) -> Any:  # noqa: ANN002, ANN003, ANN401
             # Initialize thread-local active functions set if it doesn't exist
             if not hasattr(self._thread_local, "active_functions"):
                 self._thread_local.active_functions = set()

codeflash/benchmarking/instrument_codeflash_trace.py

@@ -1,9 +1,16 @@
-from pathlib import Path
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional, Union
 
 import isort
 import libcst as cst
 
-from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from libcst import BaseStatement, ClassDef, FlattenSentinel, FunctionDef, RemovalSentinel
+
+    from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 
 
 class AddDecoratorTransformer(cst.CSTTransformer):
@@ -15,33 +22,35 @@ def __init__(self, target_functions: set[tuple[str, str]]) -> None:
         self.function_name = ""
         self.decorator = cst.Decorator(decorator=cst.Name(value="codeflash_trace"))
 
-    def leave_ClassDef(self, original_node, updated_node):
+    def leave_ClassDef(
+        self, original_node: ClassDef, updated_node: ClassDef
+    ) -> Union[BaseStatement, FlattenSentinel[BaseStatement], RemovalSentinel]:
         if self.class_name == original_node.name.value:
             self.class_name = ""  # Even if nested classes are not visited, this function is still called on them
         return updated_node
 
-    def visit_ClassDef(self, node):
+    def visit_ClassDef(self, node: ClassDef) -> Optional[bool]:
         if self.class_name:  # Don't go into nested class
             return False
-        self.class_name = node.name.value
+        self.class_name = node.name.value  # noqa: RET503
 
-    def visit_FunctionDef(self, node):
+    def visit_FunctionDef(self, node: FunctionDef) -> Optional[bool]:
         if self.function_name:  # Don't go into nested function
             return False
-        self.function_name = node.name.value
+        self.function_name = node.name.value  # noqa: RET503
 
-    def leave_FunctionDef(self, original_node, updated_node):
+    def leave_FunctionDef(self, original_node: FunctionDef, updated_node: FunctionDef) -> FunctionDef:
         if self.function_name == original_node.name.value:
             self.function_name = ""
         if (self.class_name, original_node.name.value) in self.target_functions:
             # Add the new decorator after any existing decorators, so it gets executed first
-            updated_decorators = list(updated_node.decorators) + [self.decorator]
+            updated_decorators = [*list(updated_node.decorators), self.decorator]
             self.added_codeflash_trace = True
             return updated_node.with_changes(decorators=updated_decorators)
 
         return updated_node
 
-    def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module:
+    def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module:  # noqa: ARG002
         # Create import statement for codeflash_trace
         if not self.added_codeflash_trace:
             return updated_node
@@ -68,7 +77,7 @@ def add_codeflash_decorator_to_code(code: str, functions_to_optimize: list[Funct
 
     Args:
         code: The source code as a string
-        function_to_optimize: The FunctionToOptimize instance containing function details
+        functions_to_optimize: List of FunctionToOptimize instances containing function details
 
     Returns:
         The modified source code as a string

codeflash/benchmarking/plugin/plugin.py

@@ -186,7 +186,7 @@ def get_benchmark_timings(trace_path: Path) -> dict[BenchmarkKey, int]:
 
     # Pytest hooks
     @pytest.hookimpl
-    def pytest_sessionfinish(self, session, exitstatus):
+    def pytest_sessionfinish(self, session, exitstatus) -> None:  # noqa: ANN001, ARG002
         """Execute after whole test run is completed."""
         # Write any remaining benchmark timings to the database
         codeflash_trace.close()
@@ -196,24 +196,24 @@ def pytest_sessionfinish(self, session, exitstatus):
         self.close()
 
     @staticmethod
-    def pytest_addoption(parser):
+    def pytest_addoption(parser: pytest.Parser) -> None:
         parser.addoption("--codeflash-trace", action="store_true", default=False, help="Enable CodeFlash tracing")
 
     @staticmethod
-    def pytest_plugin_registered(plugin, manager):
+    def pytest_plugin_registered(plugin, manager) -> None:  # noqa: ANN001
         # Not necessary since run with -p no:benchmark, but just in case
         if hasattr(plugin, "name") and plugin.name == "pytest-benchmark":
             manager.unregister(plugin)
 
     @staticmethod
-    def pytest_configure(config):
+    def pytest_configure(config: pytest.Config) -> None:
         """Register the benchmark marker."""
         config.addinivalue_line(
             "markers", "benchmark: mark test as a benchmark that should be run with codeflash tracing"
         )
 
     @staticmethod
-    def pytest_collection_modifyitems(config, items):
+    def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None:
         # Skip tests that don't have the benchmark fixture
         if not config.getoption("--codeflash-trace"):
             return
@@ -235,30 +235,30 @@ def pytest_collection_modifyitems(config, items):
                 item.add_marker(skip_no_benchmark)
 
     # Benchmark fixture
-    class Benchmark:
-        def __init__(self, request):
+    class Benchmark:  # noqa: D106
+        def __init__(self, request: pytest.FixtureRequest) -> None:
             self.request = request
 
-        def __call__(self, func, *args, **kwargs):
+        def __call__(self, func, *args, **kwargs):  # type: ignore  # noqa: ANN001, ANN002, ANN003, ANN204, PGH003
             """Handle both direct function calls and decorator usage."""
             if args or kwargs:
                 # Used as benchmark(func, *args, **kwargs)
                 return self._run_benchmark(func, *args, **kwargs)
 
             # Used as @benchmark decorator
-            def wrapped_func(*args, **kwargs):
+            def wrapped_func(*args, **kwargs):  # noqa: ANN002, ANN003, ANN202
                 return func(*args, **kwargs)
 
-            result = self._run_benchmark(func)
+            self._run_benchmark(func)
             return wrapped_func
 
-        def _run_benchmark(self, func, *args, **kwargs):
+        def _run_benchmark(self, func, *args, **kwargs):  # noqa: ANN001, ANN002, ANN003, ANN202
             """Actual benchmark implementation."""
             benchmark_module_path = module_name_from_file_path(
                 Path(str(self.request.node.fspath)), Path(codeflash_benchmark_plugin.project_root)
             )
             benchmark_function_name = self.request.node.name
-            line_number = int(str(sys._getframe(2).f_lineno))  # 2 frames up in the call stack
+            line_number = int(str(sys._getframe(2).f_lineno))  # 2 frames up in the call stack  # noqa: SLF001
             # Set env vars
             os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = benchmark_function_name
             os.environ["CODEFLASH_BENCHMARK_MODULE_PATH"] = benchmark_module_path
@@ -284,7 +284,7 @@ def _run_benchmark(self, func, *args, **kwargs):
 
     @staticmethod
     @pytest.fixture
-    def benchmark(request):
+    def benchmark(request: pytest.FixtureRequest) -> object:
         if not request.config.getoption("--codeflash-trace"):
             return None
 

codeflash/benchmarking/replay_test.py

@@ -47,7 +47,10 @@ def get_function_alias(module: str, function_name: str) -> str:
 
 
 def create_trace_replay_test_code(
-    trace_file: str, functions_data: list[dict[str, Any]], test_framework: str = "pytest", max_run_count=256
+    trace_file: str,
+    functions_data: list[dict[str, Any]],
+    test_framework: str = "pytest",
+    max_run_count=256,  # noqa: ANN001
 ) -> str:
     """Create a replay test for functions based on trace data.
 

codeflash/cli_cmds/cli_common.py

@@ -43,7 +43,7 @@ def inquirer_wrapper(func: Callable[..., str | bool], *args: str | bool, **kwarg
     return func(*new_args, **new_kwargs)
 
 
-def split_string_to_cli_width(string: str, is_confirm: bool = False) -> list[str]:
+def split_string_to_cli_width(string: str, is_confirm: bool = False) -> list[str]:  # noqa: FBT001, FBT002
     cli_width, _ = shutil.get_terminal_size()
     # split string to lines that accommodate "[?] " prefix
     cli_width -= len("[?] ")

codeflash/cli_cmds/cmd_init.py

@@ -34,7 +34,7 @@
     from argparse import Namespace
 
 CODEFLASH_LOGO: str = (
-    f"{LF}"
+    f"{LF}"  # noqa: ISC003
     r"                   _          ___  _               _     " + f"{LF}"
     r"                  | |        / __)| |             | |    " + f"{LF}"
     r"  ____   ___    _ | |  ____ | |__ | |  ____   ___ | | _  " + f"{LF}"
@@ -126,7 +126,8 @@ def ask_run_end_to_end_test(args: Namespace) -> None:
 
 
 def should_modify_pyproject_toml() -> bool:
-    """Check if the current directory contains a valid pyproject.toml file with codeflash config
+    """Check if the current directory contains a valid pyproject.toml file with codeflash config.
+
     If it does, ask the user if they want to re-configure it.
     """
     from rich.prompt import Confirm
@@ -144,12 +145,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:
@@ -469,7 +469,7 @@ def check_for_toml_or_setup_file() -> str | None:
     return cast("str", project_name)
 
 
-def install_github_actions(override_formatter_check: bool = False) -> None:
+def install_github_actions(override_formatter_check: bool = False) -> None:  # noqa: FBT001, FBT002
     try:
         config, config_file_path = parse_config_file(override_formatter_check=override_formatter_check)
 
@@ -566,28 +566,22 @@ def install_github_actions(override_formatter_check: bool = False) -> None:
 
 def determine_dependency_manager(pyproject_data: dict[str, Any]) -> DependencyManager:
     """Determine which dependency manager is being used based on pyproject.toml contents."""
+    result = DependencyManager.UNKNOWN
     if (Path.cwd() / "poetry.lock").exists():
-        return DependencyManager.POETRY
-    if (Path.cwd() / "uv.lock").exists():
-        return DependencyManager.UV
-    if "tool" not in pyproject_data:
-        return DependencyManager.PIP
-
-    tool_section = pyproject_data["tool"]
-
-    # Check for poetry
-    if "poetry" in tool_section:
-        return DependencyManager.POETRY
-
-    # Check for uv
-    if any(key.startswith("uv") for key in tool_section):
-        return DependencyManager.UV
-
-    # Look for pip-specific markers
-    if "pip" in tool_section or "setuptools" in tool_section:
-        return DependencyManager.PIP
-
-    return DependencyManager.UNKNOWN
+        result = DependencyManager.POETRY
+    elif (Path.cwd() / "uv.lock").exists():
+        result = DependencyManager.UV
+    elif "tool" not in pyproject_data:
+        result = DependencyManager.PIP
+    else:
+        tool_section = pyproject_data["tool"]
+        if "poetry" in tool_section:
+            result = DependencyManager.POETRY
+        elif any(key.startswith("uv") for key in tool_section):
+            result = DependencyManager.UV
+        elif "pip" in tool_section or "setuptools" in tool_section:
+            result = DependencyManager.PIP
+    return result
 
 
 def get_codeflash_github_action_command(dep_manager: DependencyManager) -> str:
@@ -642,7 +636,10 @@ def get_github_action_working_directory(toml_path: Path, git_root: Path) -> str:
 
 
 def customize_codeflash_yaml_content(
-    optimize_yml_content: str, config: tuple[dict[str, Any], Path], git_root: Path, benchmark_mode: bool = False
+    optimize_yml_content: str,
+    config: tuple[dict[str, Any], Path],
+    git_root: Path,
+    benchmark_mode: bool = False,  # noqa: FBT001, FBT002
 ) -> str:
     module_path = str(Path(config["module_root"]).relative_to(git_root) / "**")
     optimize_yml_content = optimize_yml_content.replace("{{ codeflash_module_path }}", module_path)
@@ -878,7 +875,7 @@ def test_sort(self):
         input = list(reversed(range(100)))
         output = sorter(input)
         self.assertEqual(output, list(range(100)))
-"""
+"""  # noqa: PTH119
     elif args.test_framework == "pytest":
         bubble_sort_test_content = f"""from {Path(args.module_root).name}.bubble_sort import sorter
 
@@ -959,10 +956,8 @@ def ask_for_telemetry() -> bool:
     """Prompt the user to enable or disable telemetry."""
     from rich.prompt import Confirm
 
-    enable_telemetry = Confirm.ask(
+    return Confirm.ask(
         "⚡️ Would you like to enable telemetry to help us improve the Codeflash experience?",
         default=True,
         show_default=True,
     )
-
-    return enable_telemetry

codeflash/cli_cmds/logging_config.py

@@ -27,7 +27,7 @@ def set_level(level: int, *, echo_setting: bool = True) -> None:
                 ],
                 force=True,
             )
-            logging.info("Verbose DEBUG logging enabled")
+            logging.info("Verbose DEBUG logging enabled")  # noqa: LOG015
         else:
-            logging.info("Logging level set to INFO")
+            logging.info("Logging level set to INFO")  # noqa: LOG015
     console.rule()