shifted benchmark class in plugin, improved display of benchmark info

Author: alvin-r

codeflash/benchmarking/plugin/plugin.py

@@ -5,6 +5,29 @@
 import os
 class CodeFlashBenchmarkPlugin:
     benchmark_timings = []
+
+    class Benchmark:
+        def __init__(self, request):
+            self.request = request
+
+        def __call__(self, func, *args, **kwargs):
+            benchmark_file_name = self.request.node.fspath.basename
+            benchmark_function_name = self.request.node.name
+            line_number = str(sys._getframe(1).f_lineno)  # 1 frame up in the call stack
+
+            os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = benchmark_function_name
+            os.environ["CODEFLASH_BENCHMARK_FILE_NAME"] = benchmark_file_name
+            os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = line_number
+            os.environ["CODEFLASH_BENCHMARKING"] = "True"
+
+            start = time.perf_counter_ns()
+            result = func(*args, **kwargs)
+            end = time.perf_counter_ns()
+
+            os.environ["CODEFLASH_BENCHMARKING"] = "False"
+            CodeFlashBenchmarkPlugin.benchmark_timings.append(
+                (benchmark_file_name, benchmark_function_name, line_number, end - start))
+            return result
     @staticmethod
     def pytest_addoption(parser):
         parser.addoption(
@@ -36,23 +59,4 @@ def benchmark(request):
         if not request.config.getoption("--codeflash-trace"):
             return None
 
-        class Benchmark:
-
-            def __call__(self, func, *args, **kwargs):
-                benchmark_file_name = request.node.fspath.basename
-                benchmark_function_name = request.node.name
-                line_number = str(sys._getframe(1).f_lineno)  # 1 frame up in the call stack
-                os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = benchmark_function_name
-                os.environ["CODEFLASH_BENCHMARK_FILE_NAME"] = benchmark_file_name
-                os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = line_number
-                os.environ["CODEFLASH_BENCHMARKING"] = "True"
-
-                start = time.perf_counter_ns()
-                result = func(*args, **kwargs)
-                end = time.perf_counter_ns()
-
-                os.environ["CODEFLASH_BENCHMARKING"] = "False"
-                CodeFlashBenchmarkPlugin.benchmark_timings.append((benchmark_file_name, benchmark_function_name, line_number, end - start))
-                return result
-
-        return Benchmark()
+        return CodeFlashBenchmarkPlugin.Benchmark(request)

codeflash/benchmarking/utils.py

@@ -1,17 +1,23 @@
-def print_benchmark_table(function_benchmark_timings: dict[str,dict[str,int]], total_benchmark_timings: dict[str,int]):
-    # Define column widths
-    benchmark_col_width = 50
-    time_col_width = 15
+from rich.console import Console
+from rich.table import Table
 
-    # Print table header
-    header = f"{'Benchmark Test':{benchmark_col_width}} | {'Total Time (ms)':{time_col_width}} | {'Function Time (ms)':{time_col_width}} | {'Percentage (%)':{time_col_width}}"
-    print(header)
-    print("-" * len(header))
+
+def print_benchmark_table(function_benchmark_timings: dict[str, dict[str, int]],
+                          total_benchmark_timings: dict[str, int]):
+    console = Console()
 
     # Process each function's benchmark data
     for func_path, test_times in function_benchmark_timings.items():
         function_name = func_path.split(":")[-1]
-        print(f"\n== Function: {function_name} ==")
+
+        # Create a table for this function
+        table = Table(title=f"Function: {function_name}", border_style="blue")
+
+        # Add columns
+        table.add_column("Benchmark Test", style="cyan", no_wrap=True)
+        table.add_column("Total Time (ms)", justify="right", style="green")
+        table.add_column("Function Time (ms)", justify="right", style="yellow")
+        table.add_column("Percentage (%)", justify="right", style="red")
 
         # Sort by percentage (highest first)
         sorted_tests = []
@@ -26,9 +32,16 @@ def print_benchmark_table(function_benchmark_timings: dict[str,dict[str,int]], t
 
         sorted_tests.sort(key=lambda x: x[3], reverse=True)
 
-        # Print each test's data
+        # Add rows to the table
         for test_name, total_time, func_time, percentage in sorted_tests:
             benchmark_file, benchmark_func, benchmark_line = test_name.split("::")
             benchmark_name = f"{benchmark_file}::{benchmark_func}"
-            print(f"{benchmark_name:{benchmark_col_width}} | {total_time:{time_col_width}.3f} | {func_time:{time_col_width}.3f} | {percentage:{time_col_width}.2f}")
-    print()
+            table.add_row(
+                benchmark_name,
+                f"{total_time:.3f}",
+                f"{func_time:.3f}",
+                f"{percentage:.2f}"
+            )
+
+        # Print the table
+        console.print(table)