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] 

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 "Can't do profiling or get return value from `cmd` because it is a string." 

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

80 UserWarning, 

81 ) 

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

83 # Optionally perform profiling 

84 if do_profiling: 

85 profiler = cProfile.Profile() 

86 profiler.enable() 

87 

88 try: 

89 return_value = cmd() 

90 except TypeError as e: 

91 warnings.warn( 

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

93 ) 

94 

95 if do_profiling: 

96 profiler.disable() 

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

98 

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

100 if not get_return: 

101 return_value = None 

102 

103 return FancyTimeitResult( 

104 timings=StatCounter(times), 

105 return_value=return_value, 

106 profile=profile, 

107 )