feat: support marker attributes to customize the walltime execution

Author: art049

src/pytest_codspeed/config.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import pytest
+
+
+@dataclass(frozen=True)
+class CodSpeedConfig:
+    """
+    The configuration for the codspeed plugin.
+    Usually created from the command line arguments.
+    """
+
+    warmup_time_ns: int | None = None
+    max_time_ns: int | None = None
+    max_rounds: int | None = None
+
+    @classmethod
+    def from_pytest_config(cls, config: pytest.Config) -> CodSpeedConfig:
+        warmup_time = config.getoption("--codspeed-warmup-time", None)
+        warmup_time_ns = (
+            int(warmup_time * 1_000_000_000) if warmup_time is not None else None
+        )
+        max_time = config.getoption("--codspeed-max-time", None)
+        max_time_ns = int(max_time * 1_000_000_000) if max_time is not None else None
+        return cls(
+            warmup_time_ns=warmup_time_ns,
+            max_rounds=config.getoption("--codspeed-max-rounds", None),
+            max_time_ns=max_time_ns,
+        )
+
+
+@dataclass(frozen=True)
+class BenchmarkMarkerOptions:
+    group: str | None = None
+    """The group name to use for the benchmark."""
+    min_time: int | None = None
+    """
+    The minimum time of a round (in seconds).
+    Only available in walltime mode.
+    """
+    max_time: int | None = None
+    """
+    The maximum time to run the benchmark for (in seconds).
+    Only available in walltime mode.
+    """
+    max_rounds: int | None = None
+    """
+    The maximum number of rounds to run the benchmark for.
+    Takes precedence over max_time. Only available in walltime mode.
+    """
+
+    @classmethod
+    def from_pytest_item(cls, item: pytest.Item) -> BenchmarkMarkerOptions:
+        marker = item.get_closest_marker(
+            "codspeed_benchmark"
+        ) or item.get_closest_marker("benchmark")
+        if marker is None:
+            return cls()
+        if len(marker.args) > 0:
+            raise ValueError(
+                "Positional arguments are not allowed in the benchmark marker"
+            )
+
+        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:
+            raise ValueError(
+                "Unknown kwargs passed to benchmark marker: "
+                + ", ".join(marker.kwargs.keys())
+            )
+        return options

src/pytest_codspeed/instruments/__init__.py

@@ -9,6 +9,7 @@
 
     import pytest
 
+    from pytest_codspeed.config import BenchmarkMarkerOptions
     from pytest_codspeed.plugin import CodSpeedConfig
 
     T = TypeVar("T")
@@ -27,6 +28,7 @@ def get_instrument_config_str_and_warns(self) -> tuple[str, list[str]]: ...
     @abstractmethod
     def measure(
         self,
+        marker_options: BenchmarkMarkerOptions,
         name: str,
         uri: str,
         fn: Callable[P, T],

src/pytest_codspeed/instruments/valgrind.py

@@ -13,7 +13,7 @@
     from pytest import Session
 
     from pytest_codspeed.instruments import P, T
-    from pytest_codspeed.plugin import CodSpeedConfig
+    from pytest_codspeed.plugin import BenchmarkMarkerOptions, CodSpeedConfig
 
 SUPPORTS_PERF_TRAMPOLINE = sys.version_info >= (3, 12)
 
@@ -35,7 +35,7 @@ def __init__(self, config: CodSpeedConfig) -> None:
     def get_instrument_config_str_and_warns(self) -> tuple[str, list[str]]:
         config = (
             f"mode: instrumentation, "
-            f"callgraph: {'enabled' if SUPPORTS_PERF_TRAMPOLINE  else 'not supported'}"
+            f"callgraph: {'enabled' if SUPPORTS_PERF_TRAMPOLINE else 'not supported'}"
         )
         warnings = []
         if not self.should_measure:
@@ -49,6 +49,7 @@ def get_instrument_config_str_and_warns(self) -> tuple[str, list[str]]:
 
     def measure(
         self,
+        marker_options: BenchmarkMarkerOptions,
         name: str,
         uri: str,
         fn: Callable[P, T],

src/pytest_codspeed/instruments/walltime.py

@@ -23,12 +23,12 @@
     from pytest import Session
 
     from pytest_codspeed.instruments import P, T
-    from pytest_codspeed.plugin import CodSpeedConfig
+    from pytest_codspeed.plugin import BenchmarkMarkerOptions, CodSpeedConfig
 
 DEFAULT_WARMUP_TIME_NS = 1_000_000_000
 DEFAULT_MAX_TIME_NS = 3_000_000_000
 TIMER_RESOLUTION_NS = get_clock_info("perf_counter").resolution * 1e9
-DEFAULT_MIN_ROUND_TIME_NS = TIMER_RESOLUTION_NS * 1_000_000
+DEFAULT_MIN_ROUND_TIME_NS = int(TIMER_RESOLUTION_NS * 1_000_000)
 
 IQR_OUTLIER_FACTOR = 1.5
 STDEV_OUTLIER_FACTOR = 3
@@ -42,16 +42,35 @@ class BenchmarkConfig:
     max_rounds: int | None
 
     @classmethod
-    def from_codspeed_config(cls, config: CodSpeedConfig) -> BenchmarkConfig:
+    def from_codspeed_config_and_marker_data(
+        cls, config: CodSpeedConfig, marker_data: BenchmarkMarkerOptions
+    ) -> BenchmarkConfig:
+        if marker_data.max_time is not None:
+            max_time_ns = int(marker_data.max_time * 1e9)
+        elif config.max_time_ns is not None:
+            max_time_ns = config.max_time_ns
+        else:
+            max_time_ns = DEFAULT_MAX_TIME_NS
+
+        if marker_data.max_rounds is not None:
+            max_rounds = marker_data.max_rounds
+        elif config.max_rounds is not None:
+            max_rounds = config.max_rounds
+        else:
+            max_rounds = None
+
+        if marker_data.min_time is not None:
+            min_round_time_ns = int(marker_data.min_time * 1e9)
+        else:
+            min_round_time_ns = DEFAULT_MIN_ROUND_TIME_NS
+
         return cls(
             warmup_time_ns=config.warmup_time_ns
             if config.warmup_time_ns is not None
             else DEFAULT_WARMUP_TIME_NS,
-            min_round_time_ns=DEFAULT_MIN_ROUND_TIME_NS,
-            max_time_ns=config.max_time_ns
-            if config.max_time_ns is not None
-            else DEFAULT_MAX_TIME_NS,
-            max_rounds=config.max_rounds,
+            min_round_time_ns=min_round_time_ns,
+            max_time_ns=max_time_ns,
+            max_rounds=max_rounds,
         )
 
 
@@ -231,6 +250,7 @@ def get_instrument_config_str_and_warns(self) -> tuple[str, list[str]]:
 
     def measure(
         self,
+        marker_options: BenchmarkMarkerOptions,
         name: str,
         uri: str,
         fn: Callable[P, T],
@@ -244,7 +264,9 @@ def measure(
             fn=fn,
             args=args,
             kwargs=kwargs,
-            config=BenchmarkConfig.from_codspeed_config(self.config),
+            config=BenchmarkConfig.from_codspeed_config_and_marker_data(
+                self.config, marker_options
+            ),
         )
         self.benchmarks.append(bench)
         return out

src/pytest_codspeed/plugin.py

@@ -14,6 +14,7 @@
 import pytest
 from _pytest.fixtures import FixtureManager
 
+from pytest_codspeed.config import BenchmarkMarkerOptions, CodSpeedConfig
 from pytest_codspeed.instruments import (
     MeasurementMode,
     get_instrument_from_mode,
@@ -58,8 +59,7 @@ def pytest_addoption(parser: pytest.Parser):
         action="store",
         type=float,
         help=(
-            "The time to warm up the benchmark for (in seconds), "
-            "only for walltime mode"
+            "The time to warm up the benchmark for (in seconds), only for walltime mode"
         ),
     )
     group.addoption(
@@ -82,27 +82,6 @@ def pytest_addoption(parser: pytest.Parser):
     )
 
 
-@dataclass(frozen=True)
-class CodSpeedConfig:
-    warmup_time_ns: int | None = None
-    max_time_ns: int | None = None
-    max_rounds: int | None = None
-
-    @classmethod
-    def from_pytest_config(cls, config: pytest.Config) -> CodSpeedConfig:
-        warmup_time = config.getoption("--codspeed-warmup-time", None)
-        warmup_time_ns = (
-            int(warmup_time * 1_000_000_000) if warmup_time is not None else None
-        )
-        max_time = config.getoption("--codspeed-max-time", None)
-        max_time_ns = int(max_time * 1_000_000_000) if max_time is not None else None
-        return cls(
-            warmup_time_ns=warmup_time_ns,
-            max_rounds=config.getoption("--codspeed-max-rounds", None),
-            max_time_ns=max_time_ns,
-        )
-
-
 @dataclass(unsafe_hash=True)
 class CodSpeedPlugin:
     is_codspeed_enabled: bool
@@ -254,20 +233,21 @@ def pytest_collection_modifyitems(
 
 def _measure(
     plugin: CodSpeedPlugin,
-    nodeid: str,
+    node: pytest.Item,
     config: pytest.Config,
     fn: Callable[P, T],
     *args: P.args,
     **kwargs: P.kwargs,
 ) -> T:
+    marker_options = BenchmarkMarkerOptions.from_pytest_item(node)
     random.seed(0)
     is_gc_enabled = gc.isenabled()
     if is_gc_enabled:
         gc.collect()
         gc.disable()
     try:
-        uri, name = get_git_relative_uri_and_name(nodeid, config.rootpath)
-        return plugin.instrument.measure(name, uri, fn, *args, **kwargs)
+        uri, name = get_git_relative_uri_and_name(node.nodeid, config.rootpath)
+        return plugin.instrument.measure(marker_options, name, uri, fn, *args, **kwargs)
     finally:
         # Ensure GC is re-enabled even if the test failed
         if is_gc_enabled:
@@ -276,13 +256,13 @@ def _measure(
 
 def wrap_runtest(
     plugin: CodSpeedPlugin,
-    nodeid: str,
+    node: pytest.Item,
     config: pytest.Config,
     fn: Callable[P, T],
 ) -> Callable[P, T]:
     @functools.wraps(fn)
     def wrapped(*args: P.args, **kwargs: P.kwargs) -> T:
-        return _measure(plugin, nodeid, config, fn, *args, **kwargs)
+        return _measure(plugin, node, config, fn, *args, **kwargs)
 
     return wrapped
 
@@ -299,7 +279,7 @@ def pytest_runtest_protocol(item: pytest.Item, nextitem: pytest.Item | None):
         return None
 
     # Wrap runtest and defer to default protocol
-    item.runtest = wrap_runtest(plugin, item.nodeid, item.config, item.runtest)
+    item.runtest = wrap_runtest(plugin, item, item.config, item.runtest)
     return None
 
 
@@ -343,9 +323,7 @@ def __call__(self, func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T
         config = self._request.config
         plugin = get_plugin(config)
         if plugin.is_codspeed_enabled:
-            return _measure(
-                plugin, self._request.node.nodeid, config, func, *args, **kwargs
-            )
+            return _measure(plugin, self._request.node, config, func, *args, **kwargs)
         else:
             return func(*args, **kwargs)
 

tests/test_pytest_plugin.py

@@ -220,6 +220,46 @@ def _():
     )
 
 
+def test_codspeed_marker_unexpected_args(pytester: pytest.Pytester) -> None:
+    pytester.makepyfile(
+        """
+        import pytest
+
+        @pytest.mark.codspeed_benchmark(
+            "positional_arg"
+        )
+        def test_bench():
+            pass
+        """
+    )
+    result = pytester.runpytest("--codspeed")
+    assert result.ret == 1
+    result.stdout.fnmatch_lines_random(
+        ["*ValueError: Positional arguments are not allowed in the benchmark marker*"],
+    )
+
+
+def test_codspeed_marker_unexpected_kwargs(pytester: pytest.Pytester) -> None:
+    pytester.makepyfile(
+        """
+        import pytest
+
+        @pytest.mark.codspeed_benchmark(
+            not_allowed=True
+        )
+        def test_bench():
+            pass
+        """
+    )
+    result = pytester.runpytest("--codspeed")
+    assert result.ret == 1
+    result.stdout.fnmatch_lines_random(
+        [
+            "*ValueError: Unknown kwargs passed to benchmark marker: not_allowed*",
+        ],
+    )
+
+
 def test_pytest_benchmark_extra_info(pytester: pytest.Pytester) -> None:
     """https://pytest-benchmark.readthedocs.io/en/latest/usage.html#extra-info"""
     pytester.makepyfile(