fix test

ruff reformat and fix linting

Author: misrasaurabh1

codeflash/benchmarking/codeflash_trace.py

@@ -25,6 +25,7 @@ def setup(self, trace_path: str) -> None:
         """Set up the database connection for direct writing.
 
         Args:
+        ----
             trace_path: Path to the trace database file
 
         """
@@ -52,6 +53,7 @@ def write_function_timings(self) -> None:
         """Write function call data directly to the database.
 
         Args:
+        ----
             data: List of function call data tuples to write
 
         """
@@ -94,9 +96,11 @@ def __call__(self, func: Callable) -> Callable:
         """Use as a decorator to trace function execution.
 
         Args:
+        ----
             func: The function to be decorated
 
         Returns:
+        -------
             The wrapped function
 
         """

codeflash/benchmarking/instrument_codeflash_trace.py

@@ -76,10 +76,12 @@ def add_codeflash_decorator_to_code(code: str, functions_to_optimize: list[Funct
     """Add codeflash_trace to a function.
 
     Args:
+    ----
         code: The source code as a string
         functions_to_optimize: List of FunctionToOptimize instances containing function details
 
     Returns:
+    -------
         The modified source code as a string
 
     """

codeflash/benchmarking/plugin/plugin.py

@@ -74,9 +74,11 @@ def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[Benchmark
         """Process the trace file and extract timing data for all functions.
 
         Args:
+        ----
             trace_path: Path to the trace file
 
         Returns:
+        -------
             A nested dictionary where:
             - Outer keys are module_name.qualified_name (module.class.function)
             - Inner keys are of type BenchmarkKey
@@ -132,9 +134,11 @@ def get_benchmark_timings(trace_path: Path) -> dict[BenchmarkKey, int]:
         """Extract total benchmark timings from trace files.
 
         Args:
+        ----
             trace_path: Path to the trace file
 
         Returns:
+        -------
             A dictionary mapping where:
             - Keys are of type BenchmarkKey
             - Values are total benchmark timing in milliseconds (with overhead subtracted)

codeflash/benchmarking/replay_test.py

@@ -55,12 +55,14 @@ def create_trace_replay_test_code(
     """Create a replay test for functions based on trace data.
 
     Args:
+    ----
         trace_file: Path to the SQLite database file
         functions_data: List of dictionaries with function info extracted from DB
         test_framework: 'pytest' or 'unittest'
         max_run_count: Maximum number of runs to include in the test
 
     Returns:
+    -------
         A string containing the test code
 
     """
@@ -218,12 +220,14 @@ def generate_replay_test(
     """Generate multiple replay tests from the traced function calls, grouped by benchmark.
 
     Args:
+    ----
         trace_file_path: Path to the SQLite database file
         output_dir: Directory to write the generated tests (if None, only returns the code)
         test_framework: 'pytest' or 'unittest'
         max_run_count: Maximum number of runs to include per function
 
     Returns:
+    -------
         Dictionary mapping benchmark names to generated test code
 
     """

codeflash/benchmarking/utils.py

@@ -83,11 +83,13 @@ def process_benchmark_data(
     """Process benchmark data and generate detailed benchmark information.
 
     Args:
+    ----
         replay_performance_gain: The performance gain from replay
         fto_benchmark_timings: Function to optimize benchmark timings
         total_benchmark_timings: Total benchmark timings
 
     Returns:
+    -------
         ProcessedBenchmarkInfo containing processed benchmark details
 
     """

codeflash/cli_cmds/cmd_init.py

@@ -34,7 +34,7 @@
     from argparse import Namespace
 
 CODEFLASH_LOGO: str = (
-    f"{LF}"  # noqa: ISC003
+    f"{LF}"
     r"                   _          ___  _               _     " + f"{LF}"
     r"                  | |        / __)| |             | |    " + f"{LF}"
     r"  ____   ___    _ | |  ____ | |__ | |  ____   ___ | | _  " + f"{LF}"

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")  # noqa: LOG015
+            logging.info("Verbose DEBUG logging enabled")
         else:
-            logging.info("Logging level set to INFO")  # noqa: LOG015
+            logging.info("Logging level set to INFO")
     console.rule()

codeflash/code_utils/checkpoint.py

@@ -47,6 +47,7 @@ def add_function_to_checkpoint(
         """Add a function to the checkpoint after it has been processed.
 
         Args:
+        ----
             function_fully_qualified_name: The fully qualified name of the function
             status: Status of optimization (e.g., "optimized", "failed", "skipped")
             additional_info: Any additional information to store about the function
@@ -104,7 +105,8 @@ def cleanup(self) -> None:
 def get_all_historical_functions(module_root: Path, checkpoint_dir: Path) -> dict[str, dict[str, str]]:
     """Get information about all processed functions, regardless of status.
 
-    Returns:
+    Returns
+    -------
         Dictionary mapping function names to their processing information
 
     """

codeflash/code_utils/line_profile_utils.py

@@ -24,6 +24,7 @@ def __init__(self, qualified_name: str, decorator_name: str) -> None:
         """Initialize the transformer.
 
         Args:
+        ----
             qualified_name: The fully qualified name of the function to add the decorator to (e.g., "MyClass.nested_func.target_func").
             decorator_name: The name of the decorator to add.
 
@@ -144,11 +145,13 @@ def add_decorator_to_qualified_function(module: cst.Module, qualified_name: str,
     """Add a decorator to a function with the exact qualified name in the source code.
 
     Args:
+    ----
         module: The Python source code as a CST module.
         qualified_name: The fully qualified name of the function to add the decorator to (e.g., "MyClass.nested_func.target_func").
         decorator_name: The name of the decorator to add.
 
     Returns:
+    -------
         The modified CST module.
 
     """

codeflash/context/code_context_extractor.py

@@ -3,17 +3,17 @@
 import os
 from collections import defaultdict
 from itertools import chain
-from pathlib import Path  # noqa: TC003
+from pathlib import Path
 from typing import TYPE_CHECKING
 
 import libcst as cst
-from libcst import CSTNode  # noqa: TC002
+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 encoded_tokens_len, get_qualified_name, path_belongs_to_site_packages
 from codeflash.context.unused_definition_remover import remove_unused_definitions_by_function_names
-from codeflash.discovery.functions_to_optimize import FunctionToOptimize  # noqa: TC001
+from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.models.models import (
     CodeContextType,
     CodeOptimizationContext,
@@ -150,13 +150,15 @@ def extract_code_string_context_from_files(
     imports, and combines them.
 
     Args:
+    ----
         helpers_of_fto: Dictionary mapping file paths to sets of Function Sources of function to optimize and its helpers
         helpers_of_helpers: Dictionary mapping file paths to sets of Function Sources of helpers of helper functions
         project_root_path: Root path of the project
         remove_docstrings: Whether to remove docstrings from the extracted code
         code_context_type: Type of code context to extract (READ_ONLY, READ_WRITABLE, or TESTGEN)
 
     Returns:
+    -------
         CodeString containing the extracted code context with necessary imports
 
     """  # noqa: D205
@@ -257,13 +259,15 @@ def extract_code_markdown_context_from_files(
     imports, and combines them into a structured markdown format.
 
     Args:
+    ----
         helpers_of_fto: Dictionary mapping file paths to sets of Function Sources of function to optimize and its helpers
         helpers_of_helpers: Dictionary mapping file paths to sets of Function Sources of helpers of helper functions
         project_root_path: Root path of the project
         remove_docstrings: Whether to remove docstrings from the extracted code
         code_context_type: Type of code context to extract (READ_ONLY, READ_WRITABLE, or TESTGEN)
 
     Returns:
+    -------
         CodeStringsMarkdown containing the extracted code context with necessary imports,
         formatted for inclusion in markdown
 
@@ -502,7 +506,8 @@ def prune_cst_for_read_writable_code(  # noqa: PLR0911
 ) -> tuple[cst.CSTNode | None, bool]:
     """Recursively filter the node and its children to build the read-writable codeblock. This contains nodes that lead to target functions.
 
-    Returns:
+    Returns
+    -------
         (filtered_node, found_target):
           filtered_node: The modified CST node or None if it should be removed.
           found_target: True if a target function was found in this node's subtree.
@@ -586,7 +591,8 @@ def prune_cst_for_read_only_code(  # noqa: PLR0911
 ) -> tuple[cst.CSTNode | None, bool]:
     """Recursively filter the node for read-only context.
 
-    Returns:
+    Returns
+    -------
         (filtered_node, found_target):
           filtered_node: The modified CST node or None if it should be removed.
           found_target: True if a target function was found in this node's subtree.
@@ -690,7 +696,8 @@ def prune_cst_for_testgen_code(  # noqa: PLR0911
 ) -> tuple[cst.CSTNode | None, bool]:
     """Recursively filter the node for testgen context.
 
-    Returns:
+    Returns
+    -------
         (filtered_node, found_target):
           filtered_node: The modified CST node or None if it should be removed.
           found_target: True if a target function was found in this node's subtree.