feat: support pytest-benchmark's pedantic API

Author: art049

src/pytest_codspeed/config.py

@@ -1,9 +1,17 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
-from typing import TYPE_CHECKING
+import dataclasses
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Generic, ParamSpec, TypeVar
+
+import pytest
+
+T = TypeVar("T")
+P = ParamSpec("P")
 
 if TYPE_CHECKING:
+    from typing import Any, Callable
+
     import pytest
 
 
@@ -64,17 +72,51 @@ def from_pytest_item(cls, item: pytest.Item) -> BenchmarkMarkerOptions:
             raise ValueError(
                 "Positional arguments are not allowed in the benchmark marker"
             )
+        kwargs = marker.kwargs
 
-        options = cls(
-            group=marker.kwargs.pop("group", None),
-            min_time=marker.kwargs.pop("min_time", None),
-            max_time=marker.kwargs.pop("max_time", None),
-            max_rounds=marker.kwargs.pop("max_rounds", None),
-        )
-
-        if len(marker.kwargs) > 0:
+        unknown_kwargs = set(kwargs.keys()) - {
+            field.name for field in dataclasses.fields(cls)
+        }
+        if unknown_kwargs:
             raise ValueError(
                 "Unknown kwargs passed to benchmark marker: "
-                + ", ".join(marker.kwargs.keys())
+                + ", ".join(sorted(unknown_kwargs))
+            )
+
+        return cls(**kwargs)
+
+
+@dataclass(frozen=True)
+class PedanticOptions(Generic[P, T]):
+    """Parameters for running a benchmark using the pedantic fixture API."""
+
+    target: Callable[P, T]
+    setup: Callable[[], Any | None] | None
+    teardown: Callable[P, Any | None] | None
+    rounds: int
+    warmup_rounds: int
+    iterations: int
+    args: tuple[Any, ...] = field(default_factory=tuple)
+    kwargs: dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if self.rounds < 0:
+            raise ValueError("rounds must be positive")
+        if self.warmup_rounds < 0:
+            raise ValueError("warmup_rounds must be non-negative")
+        if self.iterations <= 0:
+            raise ValueError("iterations must be positive")
+        if self.iterations > 1 and self.setup is not None:
+            raise ValueError(
+                "setup cannot be used with multiple iterations, use multiple rounds"
             )
-        return options
+
+    def setup_and_get_args_kwargs(self) -> tuple[tuple[Any, ...], dict[str, Any]]:
+        if self.setup is None:
+            return self.args, self.kwargs
+        maybe_result = self.setup(*self.args, **self.kwargs)
+        if maybe_result is not None:
+            if len(self.args) > 0 or len(self.kwargs) > 0:
+                raise ValueError("setup cannot return a value when args are provided")
+            return maybe_result
+        return self.args, self.kwargs

src/pytest_codspeed/instruments/__init__.py

@@ -9,7 +9,7 @@
 
     import pytest
 
-    from pytest_codspeed.config import BenchmarkMarkerOptions
+    from pytest_codspeed.config import BenchmarkMarkerOptions, PedanticOptions
     from pytest_codspeed.plugin import CodSpeedConfig
 
     T = TypeVar("T")
@@ -36,6 +36,15 @@ def measure(
         **kwargs: P.kwargs,
     ) -> T: ...
 
+    @abstractmethod
+    def measure_pedantic(
+        self,
+        marker_options: BenchmarkMarkerOptions,
+        pedantic_options: PedanticOptions[P, T],
+        name: str,
+        uri: str,
+    ) -> T: ...
+
     @abstractmethod
     def report(self, session: pytest.Session) -> None: ...
 

src/pytest_codspeed/instruments/valgrind/__init__.py

@@ -2,6 +2,7 @@
 
 import os
 import sys
+import warnings
 from typing import TYPE_CHECKING
 
 from pytest_codspeed import __semver_version__
@@ -13,6 +14,7 @@
 
     from pytest import Session
 
+    from pytest_codspeed.config import PedanticOptions
     from pytest_codspeed.instruments import P, T
     from pytest_codspeed.instruments.valgrind._wrapper import LibType
     from pytest_codspeed.plugin import BenchmarkMarkerOptions, CodSpeedConfig
@@ -81,8 +83,55 @@ def __codspeed_root_frame__() -> T:
             self.lib.stop_instrumentation()
             self.lib.dump_stats_at(uri.encode("ascii"))
 
+    def measure_pedantic(
+        self,
+        marker_options: BenchmarkMarkerOptions,
+        pedantic_options: PedanticOptions[P, T],
+        name: str,
+        uri: str,
+    ) -> T:
+        if pedantic_options.rounds != 1 or pedantic_options.iterations != 1:
+            warnings.warn(
+                "Valgrind instrument ignores rounds and iterations settings "
+                "in pedantic mode"
+            )
+        if self.lib is None:  # Thus should_measure is False
+            args, kwargs = pedantic_options.setup_and_get_args_kwargs()
+            out = pedantic_options.target(*args, **kwargs)
+            if pedantic_options.teardown is not None:
+                pedantic_options.teardown(*args, **kwargs)
+            return out
+
+        def __codspeed_root_frame__(*args, **kwargs) -> T:
+            return pedantic_options.target(*args, **kwargs)
+
+        # Warmup
+        warmup_rounds = max(
+            pedantic_options.warmup_rounds, 1 if SUPPORTS_PERF_TRAMPOLINE else 0
+        )
+        for _ in range(warmup_rounds):
+            args, kwargs = pedantic_options.setup_and_get_args_kwargs()
+            __codspeed_root_frame__(*args, **kwargs)
+            if pedantic_options.teardown is not None:
+                pedantic_options.teardown(*args, **kwargs)
+
+        # Compute the actual result of the function
+        args, kwargs = pedantic_options.setup_and_get_args_kwargs()
+        self.lib.zero_stats()
+        self.lib.start_instrumentation()
+        try:
+            out = __codspeed_root_frame__(*args, **kwargs)
+        finally:
+            self.lib.stop_instrumentation()
+            self.lib.dump_stats_at(uri.encode("ascii"))
+            if pedantic_options.teardown is not None:
+                pedantic_options.teardown(*args, **kwargs)
+
+        return out
+
     def report(self, session: Session) -> None:
         reporter = session.config.pluginmanager.get_plugin("terminalreporter")
+        assert reporter is not None, "terminalreporter not found"
         count_suffix = "benchmarked" if self.should_measure else "benchmark tested"
         reporter.write_sep(
             "=",

src/pytest_codspeed/instruments/walltime.py

@@ -18,6 +18,7 @@
 
     from pytest import Session
 
+    from pytest_codspeed.config import PedanticOptions
     from pytest_codspeed.instruments import P, T
     from pytest_codspeed.plugin import BenchmarkMarkerOptions, CodSpeedConfig
 
@@ -149,66 +150,6 @@ class Benchmark:
     stats: BenchmarkStats
 
 
-def run_benchmark(
-    name: str, uri: str, fn: Callable[P, T], args, kwargs, config: BenchmarkConfig
-) -> tuple[Benchmark, T]:
-    # Compute the actual result of the function
-    out = fn(*args, **kwargs)
-
-    # Warmup
-    times_per_round_ns: list[float] = []
-    warmup_start = start = perf_counter_ns()
-    while True:
-        start = perf_counter_ns()
-        fn(*args, **kwargs)
-        end = perf_counter_ns()
-        times_per_round_ns.append(end - start)
-        if end - warmup_start > config.warmup_time_ns:
-            break
-
-    # Round sizing
-    warmup_mean_ns = mean(times_per_round_ns)
-    warmup_iters = len(times_per_round_ns)
-    times_per_round_ns.clear()
-    iter_per_round = (
-        int(ceil(config.min_round_time_ns / warmup_mean_ns))
-        if warmup_mean_ns <= config.min_round_time_ns
-        else 1
-    )
-    if config.max_rounds is None:
-        round_time_ns = warmup_mean_ns * iter_per_round
-        rounds = int(config.max_time_ns / round_time_ns)
-    else:
-        rounds = config.max_rounds
-    rounds = max(1, rounds)
-
-    # Benchmark
-    iter_range = range(iter_per_round)
-    run_start = perf_counter_ns()
-    for _ in range(rounds):
-        start = perf_counter_ns()
-        for _ in iter_range:
-            fn(*args, **kwargs)
-        end = perf_counter_ns()
-        times_per_round_ns.append(end - start)
-
-        if end - run_start > config.max_time_ns:
-            # TODO: log something
-            break
-    benchmark_end = perf_counter_ns()
-    total_time = (benchmark_end - run_start) / 1e9
-
-    stats = BenchmarkStats.from_list(
-        times_per_round_ns,
-        rounds=rounds,
-        total_time=total_time,
-        iter_per_round=iter_per_round,
-        warmup_iters=warmup_iters,
-    )
-
-    return Benchmark(name=name, uri=uri, config=config, stats=stats), out
-
-
 class WallTimeInstrument(Instrument):
     instrument = "walltime"
 
@@ -228,21 +169,126 @@ def measure(
         *args: P.args,
         **kwargs: P.kwargs,
     ) -> T:
-        bench, out = run_benchmark(
-            name=name,
-            uri=uri,
-            fn=fn,
-            args=args,
-            kwargs=kwargs,
-            config=BenchmarkConfig.from_codspeed_config_and_marker_data(
-                self.config, marker_options
-            ),
+        benchmark_config = BenchmarkConfig.from_codspeed_config_and_marker_data(
+            self.config, marker_options
+        )
+
+        # Compute the actual result of the function
+        out = fn(*args, **kwargs)
+        # Warmup
+        times_per_round_ns: list[float] = []
+        warmup_start = start = perf_counter_ns()
+        while True:
+            start = perf_counter_ns()
+            fn(*args, **kwargs)
+            end = perf_counter_ns()
+            times_per_round_ns.append(end - start)
+            if end - warmup_start > benchmark_config.warmup_time_ns:
+                break
+
+        # Round sizing
+        warmup_mean_ns = mean(times_per_round_ns)
+        warmup_iters = len(times_per_round_ns)
+        times_per_round_ns.clear()
+        iter_per_round = (
+            int(ceil(benchmark_config.min_round_time_ns / warmup_mean_ns))
+            if warmup_mean_ns <= benchmark_config.min_round_time_ns
+            else 1
+        )
+        if benchmark_config.max_rounds is None:
+            round_time_ns = warmup_mean_ns * iter_per_round
+            rounds = int(benchmark_config.max_time_ns / round_time_ns)
+        else:
+            rounds = benchmark_config.max_rounds
+        rounds = max(1, rounds)
+
+        # Benchmark
+        iter_range = range(iter_per_round)
+        run_start = perf_counter_ns()
+        for _ in range(rounds):
+            start = perf_counter_ns()
+            for _ in iter_range:
+                fn(*args, **kwargs)
+            end = perf_counter_ns()
+            times_per_round_ns.append(end - start)
+
+            if end - run_start > benchmark_config.max_time_ns:
+                # TODO: log something
+                break
+        benchmark_end = perf_counter_ns()
+        total_time = (benchmark_end - run_start) / 1e9
+
+        stats = BenchmarkStats.from_list(
+            times_per_round_ns,
+            rounds=rounds,
+            total_time=total_time,
+            iter_per_round=iter_per_round,
+            warmup_iters=warmup_iters,
+        )
+
+        self.benchmarks.append(
+            Benchmark(name=name, uri=uri, config=benchmark_config, stats=stats)
+        )
+        return out
+
+    def measure_pedantic(
+        self,
+        marker_options: BenchmarkMarkerOptions,
+        pedantic_options: PedanticOptions[P, T],
+        name: str,
+        uri: str,
+    ) -> T:
+        benchmark_config = BenchmarkConfig.from_codspeed_config_and_marker_data(
+            self.config, marker_options
+        )
+
+        iter_range = range(pedantic_options.iterations)
+
+        # Warmup
+        for _ in range(pedantic_options.warmup_rounds):
+            args, kwargs = pedantic_options.setup_and_get_args_kwargs()
+            for _ in iter_range:
+                pedantic_options.target(*args, **kwargs)
+            if pedantic_options.teardown is not None:
+                pedantic_options.teardown(*args, **kwargs)
+
+        # Benchmark
+        times_per_round_ns: list[float] = []
+        benchmark_start = perf_counter_ns()
+        for _ in range(pedantic_options.rounds):
+            start = perf_counter_ns()
+            args, kwargs = pedantic_options.setup_and_get_args_kwargs()
+            for _ in iter_range:
+                pedantic_options.target(*args, **kwargs)
+            end = perf_counter_ns()
+            times_per_round_ns.append(end - start)
+            if pedantic_options.teardown is not None:
+                pedantic_options.teardown(*args, **kwargs)
+
+        benchmark_end = perf_counter_ns()
+        total_time = (benchmark_end - benchmark_start) / 1e9
+        stats = BenchmarkStats.from_list(
+            times_per_round_ns,
+            rounds=pedantic_options.rounds,
+            total_time=total_time,
+            iter_per_round=pedantic_options.iterations,
+            warmup_iters=pedantic_options.warmup_rounds,
+        )
+
+        # Compute the actual result of the function
+        args, kwargs = pedantic_options.setup_and_get_args_kwargs()
+        out = pedantic_options.target(*args, **kwargs)
+        if pedantic_options.teardown is not None:
+            pedantic_options.teardown(*args, **kwargs)
+
+        self.benchmarks.append(
+            Benchmark(name=name, uri=uri, config=benchmark_config, stats=stats)
         )
-        self.benchmarks.append(bench)
         return out
 
     def report(self, session: Session) -> None:
         reporter = session.config.pluginmanager.get_plugin("terminalreporter")
+        assert reporter is not None, "terminalreporter not found"
 
         if len(self.benchmarks) == 0:
             reporter.write_sep(