pass mypy & ruff

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/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)
 
@@ -564,7 +564,7 @@ def install_github_actions(override_formatter_check: bool = False) -> None:
         apologize_and_exit()
 
 
-def determine_dependency_manager(pyproject_data: dict[str, Any]) -> DependencyManager:
+def determine_dependency_manager(pyproject_data: dict[str, Any]) -> DependencyManager:  # noqa: PLR0911
     """Determine which dependency manager is being used based on pyproject.toml contents."""
     if (Path.cwd() / "poetry.lock").exists():
         return DependencyManager.POETRY
@@ -642,7 +642,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 +881,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 +962,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/code_utils/code_extractor.py

@@ -1,8 +1,8 @@
+# ruff: noqa: ARG002
 from __future__ import annotations
 
 import ast
-from pathlib import Path
-from typing import TYPE_CHECKING, Dict, Optional, Set
+from typing import TYPE_CHECKING, Optional
 
 import libcst as cst
 import libcst.matchers as m
@@ -11,23 +11,24 @@
 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 typing import List
+    from codeflash.models.models import FunctionSource
 
 
 class GlobalAssignmentCollector(cst.CSTVisitor):
     """Collects all global assignment statements."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
-        self.assignments: Dict[str, cst.Assign] = {}
-        self.assignment_order: List[str] = []
+        self.assignments: dict[str, cst.Assign] = {}
+        self.assignment_order: list[str] = []
         # Track scope depth to identify global assignments
         self.scope_depth = 0
         self.if_else_depth = 0
@@ -72,11 +73,11 @@ def visit_Assign(self, node: cst.Assign) -> Optional[bool]:
 class GlobalAssignmentTransformer(cst.CSTTransformer):
     """Transforms global assignments in the original file with those from the new file."""
 
-    def __init__(self, new_assignments: Dict[str, cst.Assign], new_assignment_order: List[str]):
+    def __init__(self, new_assignments: dict[str, cst.Assign], new_assignment_order: list[str]) -> None:
         super().__init__()
         self.new_assignments = new_assignments
         self.new_assignment_order = new_assignment_order
-        self.processed_assignments: Set[str] = set()
+        self.processed_assignments: set[str] = set()
         self.scope_depth = 0
         self.if_else_depth = 0
 
@@ -124,10 +125,11 @@ def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> c
         new_statements = list(updated_node.body)
 
         # Find assignments to append
-        assignments_to_append = []
-        for name in self.new_assignment_order:
-            if name not in self.processed_assignments and name in self.new_assignments:
-                assignments_to_append.append(self.new_assignments[name])
+        assignments_to_append = [
+            self.new_assignments[name]
+            for name in self.new_assignment_order
+            if name not in self.processed_assignments and name in self.new_assignments
+        ]
 
         if assignments_to_append:
             # Add a blank line before appending new assignments if needed
@@ -136,16 +138,20 @@ def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> c
                 new_statements.pop()  # Remove the Pass statement but keep the empty line
 
             # Add the new assignments
-            for assignment in assignments_to_append:
-                new_statements.append(cst.SimpleStatementLine([assignment], leading_lines=[cst.EmptyLine()]))
+            new_statements.extend(
+                [
+                    cst.SimpleStatementLine([assignment], leading_lines=[cst.EmptyLine()])
+                    for assignment in assignments_to_append
+                ]
+            )
 
         return updated_node.with_changes(body=new_statements)
 
 
 class GlobalStatementCollector(cst.CSTVisitor):
     """Visitor that collects all global statements (excluding imports and functions/classes)."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
         self.global_statements = []
         self.in_function_or_class = False
@@ -178,7 +184,7 @@ def visit_SimpleStatementLine(self, node: cst.SimpleStatementLine) -> None:
 class LastImportFinder(cst.CSTVisitor):
     """Finds the position of the last import statement in the module."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
         self.last_import_line = 0
         self.current_line = 0
@@ -193,7 +199,7 @@ def visit_SimpleStatementLine(self, node: cst.SimpleStatementLine) -> None:
 class ImportInserter(cst.CSTTransformer):
     """Transformer that inserts global statements after the last import."""
 
-    def __init__(self, global_statements: List[cst.SimpleStatementLine], last_import_line: int):
+    def __init__(self, global_statements: list[cst.SimpleStatementLine], last_import_line: int) -> None:
         super().__init__()
         self.global_statements = global_statements
         self.last_import_line = last_import_line
@@ -208,7 +214,7 @@ def leave_SimpleStatementLine(
         # If we're right after the last import and haven't inserted yet
         if self.current_line == self.last_import_line and not self.inserted:
             self.inserted = True
-            return cst.Module(body=[updated_node] + self.global_statements)
+            return cst.Module(body=[updated_node, *self.global_statements])
 
         return cst.Module(body=[updated_node])
 
@@ -222,7 +228,7 @@ def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> c
         return updated_node
 
 
-def extract_global_statements(source_code: str) -> List[cst.SimpleStatementLine]:
+def extract_global_statements(source_code: str) -> list[cst.SimpleStatementLine]:
     """Extract global statements from source code."""
     module = cst.parse_module(source_code)
     collector = GlobalStatementCollector()
@@ -285,8 +291,7 @@ def add_global_assignments(src_module_code: str, dst_module_code: str) -> str:
     transformer = GlobalAssignmentTransformer(new_collector.assignments, new_collector.assignment_order)
     transformed_module = original_module.visit(transformer)
 
-    dst_module_code = transformed_module.code
-    return dst_module_code
+    return transformed_module.code
 
 
 def add_needed_imports_from_module(
@@ -357,9 +362,10 @@ def add_needed_imports_from_module(
 
 
 def get_code(functions_to_optimize: list[FunctionToOptimize]) -> tuple[str | None, set[tuple[str, str]]]:
-    """Return the code for a function or methods in a Python module. functions_to_optimize is either a singleton
-    FunctionToOptimize instance, which represents either a function at the module level or a method of a class at the
-    module level, or it represents a list of methods of the same class.
+    """Return the code for a function or methods in a Python module.
+
+    functions_to_optimize is either a singleton FunctionToOptimize instance, which represents either a function at the
+    module level or a method of a class at the module level, or it represents a list of methods of the same class.
     """
     if (
         not functions_to_optimize
@@ -427,7 +433,7 @@ def find_target(node_list: list[ast.stmt], name_parts: tuple[str, str] | tuple[s
 
         return find_target(target.body, name_parts[1:])
 
-    with open(file_path, encoding="utf8") as file:
+    with file_path.open(encoding="utf8") as file:
         source_code: str = file.read()
     try:
         module_node: ast.Module = ast.parse(source_code)

codeflash/code_utils/code_replacer.py

@@ -82,7 +82,7 @@ def visit_ClassDef(self, node: cst.ClassDef) -> bool:
 
         return True
 
-    def leave_ClassDef(self, node: cst.ClassDef) -> None:
+    def leave_ClassDef(self, node: cst.ClassDef) -> None:  # noqa: ARG002
         if self.current_class:
             self.current_class = None
 
@@ -104,7 +104,7 @@ def __init__(
         )
         self.current_class = None
 
-    def visit_FunctionDef(self, node: cst.FunctionDef) -> bool:
+    def visit_FunctionDef(self, node: cst.FunctionDef) -> bool:  # noqa: ARG002
         return False
 
     def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef:
@@ -133,7 +133,7 @@ def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef
                 )
         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
         node = updated_node
         max_function_index = None
         class_index = None

codeflash/code_utils/shell_utils.py

@@ -1,10 +1,15 @@
+from __future__ import annotations
+
 import os
 import re
 from pathlib import Path
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
 from codeflash.code_utils.compat import LF
-from codeflash.either import Failure, Result, Success
+from codeflash.either import Failure, Success
+
+if TYPE_CHECKING:
+    from codeflash.either import Result
 
 if os.name == "nt":  # Windows
     SHELL_RC_EXPORT_PATTERN = re.compile(r"^set CODEFLASH_API_KEY=(cf-.*)$", re.MULTILINE)
@@ -17,7 +22,7 @@
 def read_api_key_from_shell_config() -> Optional[str]:
     try:
         shell_rc_path = get_shell_rc_path()
-        with open(shell_rc_path, encoding="utf8") as shell_rc:
+        with open(shell_rc_path, encoding="utf8") as shell_rc:  # noqa: PTH123
             shell_contents = shell_rc.read()
             matches = SHELL_RC_EXPORT_PATTERN.findall(shell_contents)
             return matches[-1] if matches else None
@@ -40,15 +45,14 @@ def get_api_key_export_line(api_key: str) -> str:
     return f"{SHELL_RC_EXPORT_PREFIX}{api_key}"
 
 
-def save_api_key_to_rc(api_key) -> Result[str, str]:
+def save_api_key_to_rc(api_key: str) -> Result[str, str]:
     shell_rc_path = get_shell_rc_path()
     api_key_line = get_api_key_export_line(api_key)
     try:
-        with open(shell_rc_path, "r+", encoding="utf8") as shell_file:
+        with open(shell_rc_path, "r+", encoding="utf8") as shell_file:  # noqa: PTH123
             shell_contents = shell_file.read()
-            if os.name == "nt":  # on Windows, we're writing a batch file
-                if not shell_contents:
-                    shell_contents = "@echo off"
+            if os.name == "nt" and not shell_contents:  # on windows we're writing to a batch file
+                shell_contents = "@echo off"
             existing_api_key = read_api_key_from_shell_config()
 
             if existing_api_key: