Coverage for muutils / timeit_fancy.py: 94%

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

2 

3from __future__ import annotations 

4 

5import pstats 

6import timeit 

7import cProfile 

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

9import warnings 

10 

11from muutils.statcounter import StatCounter 

12 

13T_return = TypeVar("T_return") 

14 

15 

16class FancyTimeitResult(NamedTuple): 

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

18 

19 timings: StatCounter 

20 return_value: T_return # type: ignore[valid-type] # pyright: ignore[reportGeneralTypeIssues] 

21 profile: Union[pstats.Stats, None] 

22 

23 

24def timeit_fancy( 

25 cmd: Union[Callable[[], T_return], str], 

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_return] | 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_return | None = None 

76 if (get_return or do_profiling) and isinstance(cmd, str): 

77 warnings.warn( 

78 ( 

79 "Can't do profiling or get return value from `cmd` because it is a string." 

80 + " If you want to get the return value, pass a callable instead." 

81 ), 

82 UserWarning, 

83 ) 

84 if (get_return or do_profiling) and not isinstance(cmd, str): 

85 # Optionally perform profiling 

86 if do_profiling: 

87 profiler: cProfile.Profile = cProfile.Profile() 

88 profiler.enable() 

89 

90 try: 

91 return_value = cmd() 

92 except TypeError as e: 

93 warnings.warn( 

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

95 ) 

96 

97 if do_profiling: 

98 # profiler is def bound here 

99 assert isinstance(profiler, cProfile.Profile) # pyright: ignore[reportPossiblyUnboundVariable] 

100 profiler.disable() 

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

102 

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

104 if not get_return: 

105 return_value = None 

106 

107 return FancyTimeitResult( 

108 timings=StatCounter(times), 

109 # TYPING: Argument is incorrect: Expected `typing.TypeVar`, found `None | @Todo`tyinvalid-argument-type 

110 # no idea how to fix 

111 return_value=return_value, # type: ignore[invalid-argument-type] 

112 profile=profile, 

113 )