Update walltime.py

Author: Dreamsorcerer

src/pytest_codspeed/instruments/walltime.py

@@ -2,6 +2,7 @@
 
 import os
 import warnings
+from collections.abc import Awaitable
 from dataclasses import asdict, dataclass
 from math import ceil
 from statistics import mean, quantiles, stdev
@@ -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:
         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,9 +252,48 @@ 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
+
+    async def measure_async(
+        self,
+        marker_options: BenchmarkMarkerOptions,
+        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(  # noqa: C901
+    def _measure_pedantic_iter(  # noqa: C901
         self,
         marker_options: BenchmarkMarkerOptions,
         pedantic_options: PedanticOptions[T],
@@ -273,16 +304,12 @@ def measure_pedantic(  # noqa: C901
             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[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: