Add type stubs for BenchmarkFixture

src/pytest_benchmark/fixture.pyi

@@ -0,0 +1,80 @@
+import cProfile
+import pstats
+from collections.abc import Callable
+from collections.abc import Hashable
+from typing import Any
+from typing import ParamSpec
+from typing import TypeAlias
+from typing import TypeVar
+from typing import overload
+
+from .stats import Metadata
+
+statistics: Any
+statistics_error: str | None
+
+_Args: TypeAlias = tuple[Any, ...] | tuple[()]
+_Kwargs: TypeAlias = dict[str, Any]
+_SetupFunc: TypeAlias = Callable[..., tuple[_Args, _Kwargs]]
+_TeardownFunc: TypeAlias = Callable[..., Any]
+
+_P = ParamSpec('_P')
+_R = TypeVar('_R')
+
+class BenchmarkFixture:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        self.name: str
+        self.fullname: str
+        self.disabled: bool
+        self.param: str | None
+        self.params: dict[str, Any] | None
+        self.group: str | None
+        self.has_error: bool
+        self.extra_info: dict[Hashable, Any]
+        self.skipped: bool
+
+        self.cprofile: cProfile.Profile
+        self.cprofile_loops: int | None
+        self.cprofile_dump: str | None
+
+        self.cprofile_stats: pstats.Stats | None
+        self.stats: Metadata | None
+
+    @property
+    def enabled(self) -> bool: ...
+    def __call__(  # Expose `function_to_benchmark` *args/**kwargs to `__call__`
+        self,
+        function_to_benchmark: Callable[_P, _R],
+        *args: _P.args,
+        **kwargs: _P.kwargs,
+    ) -> _R: ...
+    @overload
+    def pedantic(  # Provided `args` and/or `kwargs` (prevent `setup`)
+        self,
+        target: Callable[_P, _R],
+        args: _Args = ...,
+        kwargs: _Kwargs | None = ...,
+        *,
+        teardown: _TeardownFunc | None = ...,
+        rounds: int = ...,
+        warmup_rounds: int = ...,
+        iterations: int = ...,
+    ) -> _R: ...
+    @overload
+    def pedantic(  # Provided `setup` (prevent `args`/`kwargs`)
+        self,
+        target: Callable[_P, _R],
+        *,
+        setup: _SetupFunc | None = ...,
+        teardown: _TeardownFunc | None = ...,
+        rounds: int = ...,
+        warmup_rounds: int = ...,
+        iterations: int = ...,
+    ) -> _R: ...
+    def weave(
+        self,
+        target: Callable[_P, _R],
+        **kwargs: Any,
+    ) -> None: ...
+
+    patch = weave

src/pytest_benchmark/stats.pyi

@@ -0,0 +1,169 @@
+import pstats
+from collections.abc import Hashable
+from functools import cached_property
+from typing import Any
+from typing import Literal
+from typing import NotRequired
+from typing import TypeAlias
+from typing import TypedDict
+from typing import TypeVar
+
+from .fixture import BenchmarkFixture
+
+_D = TypeVar('_D')
+
+# Typed Dicts for as_dict returns
+
+class _StatsDict(TypedDict):
+    min: float
+    max: float
+    mean: float
+    stddev: float
+    rounds: int
+    median: float
+    iqr: float
+    q1: float
+    q3: float
+    iqr_outliers: int
+    stddev_outliers: int
+    outliers: int | str
+    ld15iqr: float
+    hd15iqr: float
+    ops: float
+    total: float
+
+class _cprofile_stats(TypedDict):
+    ncalls_recursion: str
+    ncalls: int
+    tottime: float
+    tottime_per: float
+    cumtime: float
+    cumime_per: float
+    function_name: str
+
+# Each flag for Metadata.as_dict() can change the included keys
+
+class _MetadataDict_stats(_StatsDict):
+    iterations: int
+    data: NotRequired[list[float]]
+
+class _MetadataDict(TypedDict):
+    group: str | None
+    name: str
+    fullname: str
+    params: dict[str, Any] | None
+    param: str | None
+    extra_info: dict[Hashable, Any]
+    options: dict[str, Any]
+    cprofile: NotRequired[list[_cprofile_stats]]
+    stats: NotRequired[_MetadataDict_stats]
+
+class _FlatMetadataDict(_MetadataDict, _StatsDict): ...
+
+cProfileStats = Literal[
+    'cumtime',
+    'tottime',
+    'ncalls',
+    'ncalls_recursion',
+    'tottime_per',
+    'cumtime_per',
+    'function_name',
+]
+
+Field = Literal[
+    'min',
+    'max',
+    'mean',
+    'stddev',
+    'rounds',
+    'median',
+    'iqr',
+    'q1',
+    'q3',
+    'iqr_outliers',
+    'stddev_outliers',
+    'outliers',
+    'ld15iqr',
+    'hd15iqr',
+    'ops',
+    'total',
+]
+
+cProfileFilter: TypeAlias = tuple[cProfileStats | None, int]
+
+class Stats:
+    fields: tuple[Field, ...]
+
+    def __init__(self) -> None:
+        self.data: list[float]
+
+    def __bool__(self) -> bool: ...
+    def __nonzero__(self) -> bool: ...
+    def as_dict(self) -> _StatsDict: ...
+    def update(self, duration: float) -> None: ...
+    @cached_property
+    def sorted_data(self) -> list[float]: ...
+    @cached_property
+    def total(self) -> float: ...
+    @cached_property
+    def min(self) -> float: ...
+    @cached_property
+    def max(self) -> float: ...
+    @cached_property
+    def mean(self) -> float: ...
+    @cached_property
+    def stddev(self) -> float: ...
+    @property
+    def stddev_outliers(self) -> int: ...
+    @cached_property
+    def rounds(self) -> int: ...
+    @cached_property
+    def median(self) -> float: ...
+    @cached_property
+    def ld15iqr(self) -> float: ...
+    @cached_property
+    def hd15iqr(self) -> float: ...
+    @cached_property
+    def q1(self) -> float: ...
+    @cached_property
+    def q3(self) -> float: ...
+    @cached_property
+    def iqr(self) -> float: ...
+    @cached_property
+    def iqr_outliers(self) -> int: ...
+    @cached_property
+    def outliers(self) -> str: ...
+    @cached_property
+    def ops(self) -> float: ...
+
+class Metadata:
+    def __init__(self, fixture: BenchmarkFixture, iterations: int, options: dict[str, Any]) -> None:
+        self.name: str
+        self.fullname: str
+        self.group: str | None
+        self.param: str | None
+        self.params: dict[str, Any] | None
+        self.extra_info: dict[Hashable, Any]
+        self.cprofile_stats: pstats.Stats | None
+
+        self.iterations: int
+        self.stats: Stats | None
+        self.options: dict[str, Any]
+        self.fixture: BenchmarkFixture
+
+    def __bool__(self) -> bool: ...
+    def __nonzero__(self) -> bool: ...
+    def get(self, key: str, default: _D = ...) -> Any | _D: ...
+    def __getitem__(self, key: str) -> Any: ...
+    @property
+    def has_error(self) -> bool: ...
+    def as_dict(
+        self,
+        include_data: bool = True,
+        flat: bool = False,
+        stats: bool = True,
+        cprofile: cProfileFilter | None = None,
+    ) -> _MetadataDict | _FlatMetadataDict: ...
+    def update(self, duration: float) -> None: ...
+
+def normalize_stats(stats: Stats) -> Stats: ...

tests/test_stubs.py

@@ -0,0 +1,71 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import cProfile
+    import pstats
+    from collections.abc import Hashable
+    from typing import Any
+    from typing import assert_type
+    from typing import type_check_only
+
+    from pytest_benchmark.fixture import BenchmarkFixture
+    from pytest_benchmark.stats import Metadata
+    from pytest_benchmark.stats import _FlatMetadataDict  # type: ignore
+    from pytest_benchmark.stats import _MetadataDict  # type: ignore
+
+    @type_check_only
+    def test_benchmarkfixture_hints(benchmark: BenchmarkFixture):
+        # Attributes
+        assert_type(benchmark.name, str)
+        assert_type(benchmark.fullname, str)
+        assert_type(benchmark.disabled, bool)
+        assert_type(benchmark.param, str | None)
+        assert_type(benchmark.params, dict[str, Any] | None)
+        assert_type(benchmark.group, str | None)
+        assert_type(benchmark.has_error, bool)
+        assert_type(benchmark.extra_info, dict[Hashable, Any])
+        assert_type(benchmark.skipped, bool)
+        assert_type(benchmark.cprofile, cProfile.Profile)
+        assert_type(benchmark.cprofile_loops, int | None)
+        assert_type(benchmark.cprofile_dump, str | None)
+        assert_type(benchmark.cprofile_stats, pstats.Stats | None)
+        assert_type(benchmark.stats, Metadata | None)
+
+        if benchmark.stats:
+            assert_type(benchmark.stats.as_dict(), _MetadataDict | _FlatMetadataDict)
+
+        # Properties
+        assert_type(benchmark.enabled, bool)
+
+        # Methods
+
+        # Test that types are maintained when __call__ ing the fixture
+        assert_type(benchmark.__call__(map, len, 'a'), map[int])
+        assert_type(benchmark.__call__(len, []), int)
+        assert_type(benchmark.__call__(repr, []), str)
+
+        # `args`/`kwargs`
+        assert_type(benchmark.pedantic(len, args=([],)), int)
+        assert_type(benchmark.pedantic(len, kwargs={}), int)
+        assert_type(benchmark.pedantic(len, args=(), kwargs={}), int)
+        assert_type(benchmark.pedantic(len, (), {}), int)
+        assert_type(benchmark.pedantic(len, kwargs={}, rounds=10, teardown=print), int)
+
+        # `setup`
+        assert_type(benchmark.pedantic(len, setup=lambda: ((), {})), int)
+        assert_type(benchmark.pedantic(len, setup=lambda: (([1, 2, 3],), {})), int)
+
+        # Reverse Positionals (Should show Type Error)
+        benchmark.pedantic(
+            len,
+            {},  # type: ignore[call-overload]
+            (),  # type: ignore[call-overload]
+        )
+
+        # Pass `args`/`kwargs` and `setup` (Should show Type Error)
+        benchmark.pedantic(
+            len,
+            args=(),
+            kwargs={},
+            setup=lambda: ((), {}),  # type: ignore[call-overload]
+        )