1"`timeit_fancy` is just a fancier version of timeit with more options"

2

3 from __future__ import annotations

4

5 import pstats

6 import timeit

7 import cProfile

8 from typing import Callable, Union, TypeVar, NamedTuple, Any

9 import warnings

10

11 from muutils.statcounter import StatCounter

12

13 T = TypeVar("T")

14

15

16 class FancyTimeitResult(NamedTuple):

17 """return type of `timeit_fancy`"""

18

19 timings: StatCounter

20 return_value: T # type: ignore[valid-type]

21 profile: Union[pstats.Stats, None]

22

23

24 def timeit_fancy(

25 cmd: Callable[[], T],

26 setup: Union[str, Callable[[], Any]] = lambda: None,

27 repeats: int = 5,

28 namespace: Union[dict[str, Any], None] = None,

29 get_return: bool = True,

30 do_profiling: bool = False,

31) -> FancyTimeitResult:

32 """

33 Wrapper for `timeit` to get the fastest run of a callable with more customization options.

34

35 Approximates the functionality of the %timeit magic or command line interface in a Python callable.

36

37 # Parameters

38

39 - `cmd: Callable[[], T] | str`

40 The callable to time. If a string, it will be passed to `timeit.Timer` as the `stmt` argument.

41 - `setup: str`

42 The setup code to run before `cmd`. If a string, it will be passed to `timeit.Timer` as the `setup` argument.

43 - `repeats: int`

44 The number of times to run `cmd` to get a reliable measurement.

45 - `namespace: dict[str, Any]`

46 Passed to `timeit.Timer` constructor.

47 If `cmd` or `setup` use local or global variables, they must be passed here. See `timeit` documentation for details.

48 - `get_return: bool`

49 Whether to pass the value returned from `cmd`. If True, the return value will be appended in a tuple with execution time.

50 This is for speed and convenience so that `cmd` doesn't need to be run again in the calling scope if the return values are needed.

51 (default: `False`)

52 - `do_profiling: bool`

53 Whether to return a `pstats.Stats` object in addition to the time and return value.

54 (default: `False`)

55

56 # Returns

57

58 `FancyTimeitResult`, which is a NamedTuple with the following fields:

59

60 - `time: float`

61 The time in seconds it took to run `cmd` the minimum number of times to get a reliable measurement.

62 - `return_value: T|None`

63 The return value of `cmd` if `get_return` is `True`, otherwise `None`.

64 - `profile: pstats.Stats|None`

65 A `pstats.Stats` object if `do_profiling` is `True`, otherwise `None`.

66 """

67 timer: timeit.Timer = timeit.Timer(cmd, setup, globals=namespace)

68

69 # Perform the timing

70 times: list[float] = timer.repeat(repeats, 1)

71

72 # Optionally capture the return value

73 profile: pstats.Stats | None = None

74

75 return_value: T | None = None

76 if get_return or do_profiling:

77 # Optionally perform profiling

78 if do_profiling:

79 profiler = cProfile.Profile()

80 profiler.enable()

81

82 try:

83 return_value = cmd()

84 except TypeError as e:

85 warnings.warn(

86 f"Failed to get return value from `cmd` due to error (probably passing a string). will return `return_value=None`\n{e}",

87 )

88

89 if do_profiling:

90 profiler.disable()

91 profile = pstats.Stats(profiler).strip_dirs().sort_stats("cumulative")

92

93 # reset the return value if it wasn't requested

94 if not get_return:

95 return_value = None

96

97 return FancyTimeitResult(

98 timings=StatCounter(times),

99 return_value=return_value,

100 profile=profile,

101 )