diff --git a/src/pytest_codspeed/instruments/__init__.py b/src/pytest_codspeed/instruments/__init__.py index 08dace3..5dcbbad 100644 --- a/src/pytest_codspeed/instruments/__init__.py +++ b/src/pytest_codspeed/instruments/__init__.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: + from collections.abc import Awaitable from typing import Any, Callable, ClassVar, TypeVar import pytest @@ -37,6 +38,17 @@ def measure( **kwargs: P.kwargs, ) -> T: ... + @abstractmethod + async def measure_async( + self, + marker_options: BenchmarkMarkerOptions, + name: str, + uri: str, + fn: Callable[P, Awaitable[T]], + *args: P.args, + **kwargs: P.kwargs, + ) -> T: ... + @abstractmethod def measure_pedantic( self, @@ -46,6 +58,15 @@ def measure_pedantic( uri: str, ) -> T: ... + @abstractmethod + async def measure_pedantic_async( + self, + marker_options: BenchmarkMarkerOptions, + pedantic_options: PedanticOptions[Awaitable[T]], + name: str, + uri: str, + ) -> T: ... + @abstractmethod def report(self, session: pytest.Session) -> None: ... diff --git a/src/pytest_codspeed/instruments/valgrind.py b/src/pytest_codspeed/instruments/valgrind.py index b6667f8..b9c7dbe 100644 --- a/src/pytest_codspeed/instruments/valgrind.py +++ b/src/pytest_codspeed/instruments/valgrind.py @@ -2,6 +2,7 @@ import os import warnings +from contextlib import contextmanager from typing import TYPE_CHECKING from pytest_codspeed import __semver_version__ @@ -10,6 +11,7 @@ from pytest_codspeed.utils import SUPPORTS_PERF_TRAMPOLINE if TYPE_CHECKING: + from collections.abc import Awaitable, Iterator from typing import Any, Callable from pytest import Session @@ -52,6 +54,24 @@ def get_instrument_config_str_and_warns(self) -> tuple[str, list[str]]: ) return config, warnings + @contextmanager + def _measure_context(self, uri: str): + self.benchmark_count += 1 + + if not self.instrument_hooks: + yield + return + + # Manually call the library function to avoid an extra stack frame. Also + # call the callgrind markers directly to avoid extra overhead. + self.instrument_hooks.lib.callgrind_start_instrumentation() + try: + yield + finally: + # Ensure instrumentation is stopped even if the test failed + self.instrument_hooks.lib.callgrind_stop_instrumentation() + self.instrument_hooks.set_executed_benchmark(uri) + def measure( self, marker_options: BenchmarkMarkerOptions, @@ -61,11 +81,6 @@ def measure( *args: P.args, **kwargs: P.kwargs, ) -> T: - self.benchmark_count += 1 - - if not self.instrument_hooks: - return fn(*args, **kwargs) - def __codspeed_root_frame__() -> T: return fn(*args, **kwargs) @@ -73,23 +88,34 @@ def __codspeed_root_frame__() -> T: # Warmup CPython performance map cache __codspeed_root_frame__() - # Manually call the library function to avoid an extra stack frame. Also - # call the callgrind markers directly to avoid extra overhead. - self.instrument_hooks.lib.callgrind_start_instrumentation() - try: + with self._measure_context(uri): return __codspeed_root_frame__() - finally: - # Ensure instrumentation is stopped even if the test failed - self.instrument_hooks.lib.callgrind_stop_instrumentation() - self.instrument_hooks.set_executed_benchmark(uri) - def measure_pedantic( + async def measure_async( self, marker_options: BenchmarkMarkerOptions, - pedantic_options: PedanticOptions[T], name: str, uri: str, + fn: Callable[P, Awaitable[T]], + *args: P.args, + **kwargs: P.kwargs, ) -> T: + async def __codspeed_root_frame__() -> T: + return await fn(*args, **kwargs) + + if SUPPORTS_PERF_TRAMPOLINE: + # Warmup CPython performance map cache + await __codspeed_root_frame__() + + with self._measure_context(uri): + return await __codspeed_root_frame__() + + @contextmanager + def _measure_pedantic_context( + self, + pedantic_options: PedanticOptions[Any], + uri: str, + ) -> Iterator[None]: if pedantic_options.rounds != 1 or pedantic_options.iterations != 1: warnings.warn( "Valgrind instrument ignores rounds and iterations settings " @@ -97,11 +123,29 @@ def measure_pedantic( ) if not self.instrument_hooks: args, kwargs = pedantic_options.setup_and_get_args_kwargs() - out = pedantic_options.target(*args, **kwargs) + yield if pedantic_options.teardown is not None: pedantic_options.teardown(*args, **kwargs) - return out + return + # Compute the actual result of the function + args, kwargs = pedantic_options.setup_and_get_args_kwargs() + self.instrument_hooks.lib.callgrind_start_instrumentation() + try: + yield + finally: + self.instrument_hooks.lib.callgrind_stop_instrumentation() + self.instrument_hooks.set_executed_benchmark(uri) + if pedantic_options.teardown is not None: + pedantic_options.teardown(*args, **kwargs) + + def measure_pedantic( + self, + marker_options: BenchmarkMarkerOptions, + pedantic_options: PedanticOptions[T], + name: str, + uri: str, + ) -> T: def __codspeed_root_frame__(*args, **kwargs) -> T: return pedantic_options.target(*args, **kwargs) @@ -115,18 +159,31 @@ def __codspeed_root_frame__(*args, **kwargs) -> T: 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.instrument_hooks.lib.callgrind_start_instrumentation() - try: - out = __codspeed_root_frame__(*args, **kwargs) - finally: - self.instrument_hooks.lib.callgrind_stop_instrumentation() - self.instrument_hooks.set_executed_benchmark(uri) + with self._measure_pedantic_context(pedantic_options, uri): + return __codspeed_root_frame__(*args, **kwargs) + + async def measure_pedantic_async( + self, + marker_options: BenchmarkMarkerOptions, + pedantic_options: PedanticOptions[Awaitable[T]], + name: str, + uri: str, + ) -> T: + async def __codspeed_root_frame__(*args, **kwargs) -> T: + return await 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() + await __codspeed_root_frame__(*args, **kwargs) if pedantic_options.teardown is not None: pedantic_options.teardown(*args, **kwargs) - return out + with self._measure_pedantic_context(pedantic_options, uri): + return await __codspeed_root_frame__(*args, **kwargs) def report(self, session: Session) -> None: reporter = session.config.pluginmanager.get_plugin("terminalreporter") diff --git a/src/pytest_codspeed/instruments/walltime.py b/src/pytest_codspeed/instruments/walltime.py index f85f857..7422fd4 100644 --- a/src/pytest_codspeed/instruments/walltime.py +++ b/src/pytest_codspeed/instruments/walltime.py @@ -19,6 +19,7 @@ from pytest_codspeed.utils import SUPPORTS_PERF_TRAMPOLINE if TYPE_CHECKING: + from collections.abc import Awaitable, Iterator from typing import Any, Callable from pytest import Session @@ -182,31 +183,22 @@ def get_instrument_config_str_and_warns(self) -> tuple[str, list[str]]: ) return config_str, [] - def measure( + def _measure_iter( self, marker_options: BenchmarkMarkerOptions, name: str, uri: str, - fn: Callable[P, T], - *args: P.args, - **kwargs: P.kwargs, - ) -> T: + ) -> Iterator[None]: benchmark_config = BenchmarkConfig.from_codspeed_config_and_marker_data( self.config, marker_options ) - def __codspeed_root_frame__() -> T: - return fn(*args, **kwargs) - - # Compute the actual result of the function - out = __codspeed_root_frame__() - # Warmup times_per_round_ns: list[float] = [] warmup_start = start = perf_counter_ns() while True: start = perf_counter_ns() - __codspeed_root_frame__() + yield end = perf_counter_ns() times_per_round_ns.append(end - start) if end - warmup_start > benchmark_config.warmup_time_ns: @@ -236,7 +228,7 @@ def __codspeed_root_frame__() -> T: for _ in range(rounds): start = perf_counter_ns() for _ in iter_range: - __codspeed_root_frame__() + yield end = perf_counter_ns() times_per_round_ns.append(end - start) @@ -260,29 +252,64 @@ def __codspeed_root_frame__() -> T: self.benchmarks.append( Benchmark(name=name, uri=uri, config=benchmark_config, stats=stats) ) + + def measure( + self, + marker_options: BenchmarkMarkerOptions, + name: str, + uri: str, + fn: Callable[P, T], + *args: P.args, + **kwargs: P.kwargs, + ) -> T: + def __codspeed_root_frame__() -> T: + return fn(*args, **kwargs) + + # Compute the actual result of the function + out = __codspeed_root_frame__() + + for _ in self._measure_iter(marker_options, name, uri): + __codspeed_root_frame__() + return out - def measure_pedantic( # noqa: C901 + async def measure_async( self, marker_options: BenchmarkMarkerOptions, - pedantic_options: PedanticOptions[T], name: str, uri: str, + fn: Callable[P, Awaitable[T]], + *args: P.args, + **kwargs: P.kwargs, ) -> T: + async def __codspeed_root_frame__() -> T: + return await fn(*args, **kwargs) + + # Compute the actual result of the function + out = await __codspeed_root_frame__() + + for _ in self._measure_iter(marker_options, name, uri): + await __codspeed_root_frame__() + + return out + + def _measure_pedantic_iter( # noqa: C901 + self, + marker_options: BenchmarkMarkerOptions, + pedantic_options: PedanticOptions[Any], + name: str, + uri: str, + ) -> Iterator[tuple[range, tuple[Any, ...], dict[str, Any]]]: benchmark_config = BenchmarkConfig.from_codspeed_config_and_marker_data( self.config, marker_options ) - def __codspeed_root_frame__(*args, **kwargs) -> T: - return pedantic_options.target(*args, **kwargs) - 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: - __codspeed_root_frame__(*args, **kwargs) + yield iter_range, args, kwargs if pedantic_options.teardown is not None: pedantic_options.teardown(*args, **kwargs) @@ -294,8 +321,7 @@ def __codspeed_root_frame__(*args, **kwargs) -> T: for _ in range(pedantic_options.rounds): start = perf_counter_ns() args, kwargs = pedantic_options.setup_and_get_args_kwargs() - for _ in iter_range: - __codspeed_root_frame__(*args, **kwargs) + yield iter_range, args, kwargs end = perf_counter_ns() times_per_round_ns.append(end - start) if pedantic_options.teardown is not None: @@ -313,15 +339,56 @@ def __codspeed_root_frame__(*args, **kwargs) -> T: warmup_iters=pedantic_options.warmup_rounds, ) + self.benchmarks.append( + Benchmark(name=name, uri=uri, config=benchmark_config, stats=stats) + ) + + def measure_pedantic( + self, + marker_options: BenchmarkMarkerOptions, + pedantic_options: PedanticOptions[T], + name: str, + uri: str, + ) -> T: + def __codspeed_root_frame__(*args, **kwargs) -> T: + return pedantic_options.target(*args, **kwargs) + + for i, args, kwargs in self._measure_pedantic_iter( + marker_options, pedantic_options, name, uri + ): + for _ in i: + __codspeed_root_frame__(*args, **kwargs) + # Compute the actual result of the function args, kwargs = pedantic_options.setup_and_get_args_kwargs() out = __codspeed_root_frame__(*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) - ) + return out + + async def measure_pedantic_async( + self, + marker_options: BenchmarkMarkerOptions, + pedantic_options: PedanticOptions[Awaitable[T]], + name: str, + uri: str, + ) -> T: + async def __codspeed_root_frame__(*args, **kwargs) -> T: + return await pedantic_options.target(*args, **kwargs) + + for i, args, kwargs in self._measure_pedantic_iter( + marker_options, pedantic_options, name, uri + ): + for _ in i: + await __codspeed_root_frame__(*args, **kwargs) + + # Compute the actual result of the function + args, kwargs = pedantic_options.setup_and_get_args_kwargs() + out = await __codspeed_root_frame__(*args, **kwargs) + if pedantic_options.teardown is not None: + pedantic_options.teardown(*args, **kwargs) + return out def report(self, session: Session) -> None: diff --git a/src/pytest_codspeed/plugin.py b/src/pytest_codspeed/plugin.py index 0d2853b..3da71b8 100644 --- a/src/pytest_codspeed/plugin.py +++ b/src/pytest_codspeed/plugin.py @@ -5,7 +5,9 @@ import json import os import random +from contextlib import contextmanager from dataclasses import dataclass, field +from inspect import iscoroutinefunction from pathlib import Path from time import time from typing import TYPE_CHECKING, cast @@ -30,6 +32,7 @@ from . import __version__ if TYPE_CHECKING: + from collections.abc import Awaitable, Iterator from typing import Any, Callable, ParamSpec, TypeVar from pytest_codspeed.instruments import Instrument @@ -232,35 +235,68 @@ def pytest_collection_modifyitems( items[:] = selected -def _measure( - plugin: CodSpeedPlugin, - node: pytest.Item, - config: pytest.Config, - pedantic_options: PedanticOptions | None, - fn: Callable[..., T], - args: tuple[Any, ...], - kwargs: dict[str, Any], -) -> T: - marker_options = BenchmarkMarkerOptions.from_pytest_item(node) +@contextmanager +def _measure_context() -> Iterator[None]: 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(node.nodeid, config.rootpath) + yield + finally: + # Ensure GC is re-enabled even if the test failed + if is_gc_enabled: + gc.enable() + + +async def _async_measure( + plugin: CodSpeedPlugin, + marker_options: BenchmarkMarkerOptions, + pedantic_options: PedanticOptions | None, + name: str, + uri: str, + fn: Callable[..., Awaitable[T]], + args: tuple[Any, ...], + kwargs: dict[str, Any], +) -> T: + with _measure_context(): if pedantic_options is None: - return plugin.instrument.measure( + return await plugin.instrument.measure_async( marker_options, name, uri, fn, *args, **kwargs ) else: - return plugin.instrument.measure_pedantic( + return await plugin.instrument.measure_pedantic_async( marker_options, pedantic_options, name, uri ) - finally: - # Ensure GC is re-enabled even if the test failed - if is_gc_enabled: - gc.enable() + + +def _measure( + plugin: CodSpeedPlugin, + node: pytest.Item, + config: pytest.Config, + pedantic_options: PedanticOptions | None, + fn: Callable[..., T], + args: tuple[Any, ...], + kwargs: dict[str, Any], +) -> T: + marker_options = BenchmarkMarkerOptions.from_pytest_item(node) + uri, name = get_git_relative_uri_and_name(node.nodeid, config.rootpath) + if iscoroutinefunction(fn): + return _async_measure( # type: ignore[return-value] + plugin, marker_options, pedantic_options, name, uri, fn, args, kwargs + ) + else: + with _measure_context(): + if pedantic_options is None: + return plugin.instrument.measure( + marker_options, name, uri, fn, *args, **kwargs + ) + else: + return plugin.instrument.measure_pedantic( + marker_options, pedantic_options, name, uri + ) def wrap_runtest(