ruff format benchmarking

Author: KRRT7

codeflash/benchmarking/codeflash_trace.py

@@ -69,7 +69,7 @@ def write_function_timings(self) -> None:
                 "(function_name, class_name, module_name, file_path, benchmark_function_name, "
                 "benchmark_module_path, benchmark_line_number, function_time_ns, overhead_time_ns, args, kwargs) "
                 "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
-                self.function_calls_data
+                self.function_calls_data,
             )
             self._connection.commit()
             self.function_calls_data = []
@@ -100,7 +100,8 @@ def __call__(self, func: Callable) -> Callable:
             The wrapped function
 
         """
-        func_id = (func.__module__,func.__name__)
+        func_id = (func.__module__, func.__name__)
+
         @functools.wraps(func)
         def wrapper(*args, **kwargs):
             # Initialize thread-local active functions set if it doesn't exist
@@ -139,9 +140,19 @@ def wrapper(*args, **kwargs):
                 self._thread_local.active_functions.remove(func_id)
                 overhead_time = time.thread_time_ns() - end_time
                 self.function_calls_data.append(
-                    (func.__name__, class_name, func.__module__, func.__code__.co_filename,
-                     benchmark_function_name, benchmark_module_path, benchmark_line_number, execution_time,
-                     overhead_time, None, None)
+                    (
+                        func.__name__,
+                        class_name,
+                        func.__module__,
+                        func.__code__.co_filename,
+                        benchmark_function_name,
+                        benchmark_module_path,
+                        benchmark_line_number,
+                        execution_time,
+                        overhead_time,
+                        None,
+                        None,
+                    )
                 )
                 return result
 
@@ -155,9 +166,19 @@ def wrapper(*args, **kwargs):
                 self._thread_local.active_functions.remove(func_id)
                 overhead_time = time.thread_time_ns() - end_time
                 self.function_calls_data.append(
-                    (func.__name__, class_name, func.__module__, func.__code__.co_filename,
-                     benchmark_function_name, benchmark_module_path, benchmark_line_number, execution_time,
-                     overhead_time, None, None)
+                    (
+                        func.__name__,
+                        class_name,
+                        func.__module__,
+                        func.__code__.co_filename,
+                        benchmark_function_name,
+                        benchmark_module_path,
+                        benchmark_line_number,
+                        execution_time,
+                        overhead_time,
+                        None,
+                        None,
+                    )
                 )
                 return result
             # Flush to database every 100 calls
@@ -168,12 +189,24 @@ def wrapper(*args, **kwargs):
             self._thread_local.active_functions.remove(func_id)
             overhead_time = time.thread_time_ns() - end_time
             self.function_calls_data.append(
-                (func.__name__, class_name, func.__module__, func.__code__.co_filename,
-                 benchmark_function_name, benchmark_module_path, benchmark_line_number, execution_time,
-                 overhead_time, pickled_args, pickled_kwargs)
+                (
+                    func.__name__,
+                    class_name,
+                    func.__module__,
+                    func.__code__.co_filename,
+                    benchmark_function_name,
+                    benchmark_module_path,
+                    benchmark_line_number,
+                    execution_time,
+                    overhead_time,
+                    pickled_args,
+                    pickled_kwargs,
+                )
             )
             return result
+
         return wrapper
 
+
 # Create a singleton instance
 codeflash_trace = CodeflashTrace()

codeflash/benchmarking/instrument_codeflash_trace.py

@@ -13,57 +13,46 @@ def __init__(self, target_functions: set[tuple[str, str]]) -> None:
         self.added_codeflash_trace = False
         self.class_name = ""
         self.function_name = ""
-        self.decorator = cst.Decorator(
-            decorator=cst.Name(value="codeflash_trace")
-        )
+        self.decorator = cst.Decorator(decorator=cst.Name(value="codeflash_trace"))
 
-    def leave_ClassDef(self, original_node, updated_node):
+    def leave_ClassDef(self, original_node, updated_node):  # noqa: ANN001, ANN201, N802
         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
+            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):
-        if self.class_name: # Don't go into nested class
+    def visit_ClassDef(self, node):  # noqa: ANN001, ANN201, N802
+        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):
-        if self.function_name: # Don't go into nested function
+    def visit_FunctionDef(self, node):  # noqa: ANN001, ANN201, N802
+        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, updated_node):  # noqa: ANN001, ANN201, N802
         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.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, N802
         # Create import statement for codeflash_trace
         if not self.added_codeflash_trace:
             return updated_node
         import_stmt = cst.SimpleStatementLine(
             body=[
                 cst.ImportFrom(
                     module=cst.Attribute(
-                        value=cst.Attribute(
-                            value=cst.Name(value="codeflash"),
-                            attr=cst.Name(value="benchmarking")
-                        ),
-                        attr=cst.Name(value="codeflash_trace")
+                        value=cst.Attribute(value=cst.Name(value="codeflash"), attr=cst.Name(value="benchmarking")),
+                        attr=cst.Name(value="codeflash_trace"),
                     ),
-                    names=[
-                        cst.ImportAlias(
-                            name=cst.Name(value="codeflash_trace")
-                        )
-                    ]
+                    names=[cst.ImportAlias(name=cst.Name(value="codeflash_trace"))],
                 )
             ]
         )
@@ -73,12 +62,13 @@ def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> c
 
         return updated_node.with_changes(body=new_body)
 
+
 def add_codeflash_decorator_to_code(code: str, functions_to_optimize: list[FunctionToOptimize]) -> str:
     """Add codeflash_trace to a function.
 
     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
@@ -91,25 +81,18 @@ def add_codeflash_decorator_to_code(code: str, functions_to_optimize: list[Funct
             class_name = function_to_optimize.parents[0].name
         target_functions.add((class_name, function_to_optimize.function_name))
 
-    transformer = AddDecoratorTransformer(
-        target_functions = target_functions,
-    )
+    transformer = AddDecoratorTransformer(target_functions=target_functions)
 
     module = cst.parse_module(code)
     modified_module = module.visit(transformer)
     return modified_module.code
 
 
-def instrument_codeflash_trace_decorator(
-        file_to_funcs_to_optimize: dict[Path, list[FunctionToOptimize]]
-) -> None:
+def instrument_codeflash_trace_decorator(file_to_funcs_to_optimize: dict[Path, list[FunctionToOptimize]]) -> None:
     """Instrument codeflash_trace decorator to functions to optimize."""
     for file_path, functions_to_optimize in file_to_funcs_to_optimize.items():
         original_code = file_path.read_text(encoding="utf-8")
-        new_code = add_codeflash_decorator_to_code(
-            original_code,
-            functions_to_optimize
-        )
+        new_code = add_codeflash_decorator_to_code(original_code, functions_to_optimize)
         # Modify the code
         modified_code = isort.code(code=new_code, float_to_top=True)
 

codeflash/benchmarking/plugin/plugin.py

@@ -20,7 +20,7 @@ def __init__(self) -> None:
         self.project_root = None
         self.benchmark_timings = []
 
-    def setup(self, trace_path:str, project_root:str) -> None:
+    def setup(self, trace_path: str, project_root: str) -> None:
         try:
             # Open connection
             self.project_root = project_root
@@ -35,7 +35,7 @@ def setup(self, trace_path:str, project_root:str) -> None:
                 "benchmark_time_ns INTEGER)"
             )
             self._connection.commit()
-            self.close() # Reopen only at the end of pytest session
+            self.close()  # Reopen only at the end of pytest session
         except Exception as e:
             print(f"Database setup error: {e}")
             if self._connection:
@@ -55,14 +55,15 @@ def write_benchmark_timings(self) -> None:
             # Insert data into the benchmark_timings table
             cur.executemany(
                 "INSERT INTO benchmark_timings (benchmark_module_path, benchmark_function_name, benchmark_line_number, benchmark_time_ns) VALUES (?, ?, ?, ?)",
-                self.benchmark_timings
+                self.benchmark_timings,
             )
             self._connection.commit()
-            self.benchmark_timings = [] # Clear the benchmark timings list
+            self.benchmark_timings = []  # Clear the benchmark timings list
         except Exception as e:
             print(f"Error writing to benchmark timings database: {e}")
             self._connection.rollback()
             raise
+
     def close(self) -> None:
         if self._connection:
             self._connection.close()
@@ -196,12 +197,7 @@ def pytest_sessionfinish(self, session, exitstatus):
 
     @staticmethod
     def pytest_addoption(parser):
-        parser.addoption(
-            "--codeflash-trace",
-            action="store_true",
-            default=False,
-            help="Enable CodeFlash tracing"
-        )
+        parser.addoption("--codeflash-trace", action="store_true", default=False, help="Enable CodeFlash tracing")
 
     @staticmethod
     def pytest_plugin_registered(plugin, manager):
@@ -213,9 +209,9 @@ def pytest_plugin_registered(plugin, manager):
     def pytest_configure(config):
         """Register the benchmark marker."""
         config.addinivalue_line(
-            "markers",
-            "benchmark: mark test as a benchmark that should be run with codeflash tracing"
+            "markers", "benchmark: mark test as a benchmark that should be run with codeflash tracing"
         )
+
     @staticmethod
     def pytest_collection_modifyitems(config, items):
         # Skip tests that don't have the benchmark fixture
@@ -248,16 +244,18 @@ def __call__(self, func, *args, **kwargs):
             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):
-                return func(*args, **kwargs)
-            result = self._run_benchmark(func)
+            def wrapped_func(*inner_args, **inner_kwargs):
+                return self._run_benchmark(func, *inner_args, **inner_kwargs)
+
             return wrapped_func
 
         def _run_benchmark(self, func, *args, **kwargs):
             """Actual benchmark implementation."""
-            benchmark_module_path = module_name_from_file_path(Path(str(self.request.node.fspath)),
-                                                               Path(codeflash_benchmark_plugin.project_root))
+            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
             # Set env vars
@@ -278,7 +276,8 @@ def _run_benchmark(self, func, *args, **kwargs):
             codeflash_trace.function_call_count = 0
             # Add to the benchmark timings buffer
             codeflash_benchmark_plugin.benchmark_timings.append(
-                (benchmark_module_path, benchmark_function_name, line_number, end - start))
+                (benchmark_module_path, benchmark_function_name, line_number, end - start)
+            )
 
             return result
 
@@ -290,4 +289,5 @@ def benchmark(request):
 
         return CodeFlashBenchmarkPlugin.Benchmark(request)
 
+
 codeflash_benchmark_plugin = CodeFlashBenchmarkPlugin()

codeflash/benchmarking/pytest_new_process_trace_benchmarks.py

@@ -16,8 +16,20 @@
         codeflash_benchmark_plugin.setup(trace_file, project_root)
         codeflash_trace.setup(trace_file)
         exitcode = pytest.main(
-            [benchmarks_root, "--codeflash-trace", "-p", "no:benchmark","-p", "no:codspeed","-p", "no:cov-s", "-o", "addopts="], plugins=[codeflash_benchmark_plugin]
-        ) # Errors will be printed to stdout, not stderr
+            [
+                benchmarks_root,
+                "--codeflash-trace",
+                "-p",
+                "no:benchmark",
+                "-p",
+                "no:codspeed",
+                "-p",
+                "no:cov-s",
+                "-o",
+                "addopts=",
+            ],
+            plugins=[codeflash_benchmark_plugin],
+        )  # Errors will be printed to stdout, not stderr
 
     except Exception as e:
         print(f"Failed to collect tests: {e!s}", file=sys.stderr)