Merge pull request #34 from codeflash-ai/ruff-fixes

ruff check & format

Author: KRRT7

codeflash/api/aiservice.py

@@ -196,10 +196,9 @@ def generate_regression_tests(
         - Dict[str, str] | None: The generated regression tests and instrumented tests, or None if an error occurred.
 
         """
-        assert test_framework in [
-            "pytest",
-            "unittest",
-        ], f"Invalid test framework, got {test_framework} but expected 'pytest' or 'unittest'"
+        assert test_framework in ["pytest", "unittest"], (
+            f"Invalid test framework, got {test_framework} but expected 'pytest' or 'unittest'"
+        )
         payload = {
             "source_code_being_tested": source_code_being_tested,
             "function_to_optimize": function_to_optimize,

codeflash/cli_cmds/cmd_init.py

@@ -462,7 +462,7 @@ def determine_dependency_manager(pyproject_data: dict[str, Any]) -> DependencyMa
         return DependencyManager.POETRY
 
     # Check for uv
-    if any(key.startswith("uv") for key in tool_section.keys()):
+    if any(key.startswith("uv") for key in tool_section):
         return DependencyManager.UV
 
     # Look for pip-specific markers
@@ -555,9 +555,7 @@ def customize_codeflash_yaml_content(
 
     # Add codeflash command
     codeflash_cmd = get_codeflash_github_action_command(dep_manager)
-    optimize_yml_content = optimize_yml_content.replace("{{ codeflash_command }}", codeflash_cmd)
-
-    return optimize_yml_content
+    return optimize_yml_content.replace("{{ codeflash_command }}", codeflash_cmd)
 
 
 # Create or update the pyproject.toml file with the Codeflash dependency & configuration
@@ -596,8 +594,8 @@ def configure_pyproject_toml(setup_info: SetupInfo) -> None:
         formatter_cmds.append("disabled")
     if formatter in ["black", "ruff"]:
         try:
-            result = subprocess.run([formatter], capture_output=True, check=False)
-        except FileNotFoundError as e:
+            subprocess.run([formatter], capture_output=True, check=False)
+        except FileNotFoundError:
             click.echo(f"⚠️ Formatter not found: {formatter}, please ensure it is installed")
     codeflash_section["formatter-cmds"] = formatter_cmds
     # Add the 'codeflash' section, ensuring 'tool' section exists

codeflash/cli_cmds/console.py

@@ -3,7 +3,7 @@
 import logging
 from contextlib import contextmanager
 from itertools import cycle
-from typing import TYPE_CHECKING, Generator
+from typing import TYPE_CHECKING
 
 from rich.console import Console
 from rich.logging import RichHandler
@@ -13,6 +13,8 @@
 from codeflash.cli_cmds.logging_config import BARE_LOGGING_FORMAT
 
 if TYPE_CHECKING:
+    from collections.abc import Generator
+
     from rich.progress import TaskID
 
 DEBUG_MODE = logging.getLogger().getEffectiveLevel() == logging.DEBUG

codeflash/code_utils/code_extractor.py

@@ -115,8 +115,7 @@ def get_code(functions_to_optimize: list[FunctionToOptimize]) -> tuple[str | Non
         or (functions_to_optimize[0].parents and functions_to_optimize[0].parents[0].type != "ClassDef")
         or (
             len(functions_to_optimize[0].parents) > 1
-            or (len(functions_to_optimize) > 1)
-            and len({fn.parents[0] for fn in functions_to_optimize}) != 1
+            or ((len(functions_to_optimize) > 1) and len({fn.parents[0] for fn in functions_to_optimize}) != 1)
         )
     ):
         return None, set()

codeflash/code_utils/code_replacer.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import ast
-import re
 from collections import defaultdict
 from functools import lru_cache
 from typing import TYPE_CHECKING, Optional, TypeVar
@@ -91,10 +90,10 @@ def leave_ClassDef(self, node: cst.ClassDef) -> None:
 class OptimFunctionReplacer(cst.CSTTransformer):
     def __init__(
         self,
-        modified_functions: dict[tuple[str | None, str], cst.FunctionDef] = None,
-        new_functions: list[cst.FunctionDef] = None,
-        new_class_functions: dict[str, list[cst.FunctionDef]] = None,
-        modified_init_functions: dict[str, cst.FunctionDef] = None,
+        modified_functions: Optional[dict[tuple[str | None, str], cst.FunctionDef]] = None,
+        new_functions: Optional[list[cst.FunctionDef]] = None,
+        new_class_functions: Optional[dict[str, list[cst.FunctionDef]]] = None,
+        modified_init_functions: Optional[dict[str, cst.FunctionDef]] = None,
     ) -> None:
         super().__init__()
         self.modified_functions = modified_functions if modified_functions is not None else {}

codeflash/code_utils/code_utils.py

@@ -12,7 +12,8 @@
 
 def get_qualified_name(module_name: str, full_qualified_name: str) -> str:
     if not full_qualified_name:
-        raise ValueError("full_qualified_name cannot be empty")
+        msg = "full_qualified_name cannot be empty"
+        raise ValueError(msg)
     if not full_qualified_name.startswith(module_name):
         msg = f"{full_qualified_name} does not start with {module_name}"
         raise ValueError(msg)
@@ -46,9 +47,9 @@ def file_name_from_test_module_name(test_module_name: str, base_dir: Path) -> Pa
 def get_imports_from_file(
     file_path: Path | None = None, file_string: str | None = None, file_ast: ast.AST | None = None
 ) -> list[ast.Import | ast.ImportFrom]:
-    assert (
-        sum([file_path is not None, file_string is not None, file_ast is not None]) == 1
-    ), "Must provide exactly one of file_path, file_string, or file_ast"
+    assert sum([file_path is not None, file_string is not None, file_ast is not None]) == 1, (
+        "Must provide exactly one of file_path, file_string, or file_ast"
+    )
     if file_path:
         with file_path.open(encoding="utf8") as file:
             file_string = file.read()

codeflash/code_utils/compat.py

@@ -11,4 +11,4 @@
 
 SAFE_SYS_EXECUTABLE: str = Path(sys.executable).as_posix()
 
-IS_POSIX = os.name != "nt"
+IS_POSIX = os.name != "nt"

codeflash/code_utils/concolic_utils.py

@@ -62,7 +62,7 @@ def _split_top_level_args(self, args_str: str) -> list[str]:
 
         return result
 
-    def __init__(self):
+    def __init__(self) -> None:
         # Pre-compiling regular expressions for faster execution
         self.assert_re = re.compile(r"\s*assert\s+(.*?)(?:\s*==\s*.*)?$")
         self.unittest_re = re.compile(r"(\s*)self\.assert([A-Za-z]+)\((.*)\)$")

codeflash/code_utils/config_parser.py

@@ -81,10 +81,9 @@ def parse_config_file(config_file_path: Path | None = None) -> tuple[dict[str, A
         else:  # Default to empty list
             config[key] = []
 
-    assert config["test-framework"] in [
-        "pytest",
-        "unittest",
-    ], "In pyproject.toml, Codeflash only supports the 'test-framework' as pytest and unittest."
+    assert config["test-framework"] in ["pytest", "unittest"], (
+        "In pyproject.toml, Codeflash only supports the 'test-framework' as pytest and unittest."
+    )
     if len(config["formatter-cmds"]) > 0:
         assert config["formatter-cmds"][0] != "your-formatter $file", (
             "The formatter command is not set correctly in pyproject.toml. Please set the "

codeflash/code_utils/formatter.py

@@ -3,7 +3,6 @@
 import os
 import shlex
 import subprocess
-import sys
 from typing import TYPE_CHECKING
 
 import isort