Coverage for mlair/helpers/time_tracking.py: 97%

1"""Track time either as decorator or explicit.""" 

2import datetime as dt 

3import logging 

4import math 

5import time 

6import types 

7from functools import wraps 

8from typing import Optional 

9 

10 

11class TimeTrackingWrapper: 

12 r""" 

13 Wrapper implementation of TimeTracking class. 

14 

15 Use this implementation easily as decorator for functions, classes and class methods. Implement a custom function 

16 and decorate it for automatic time measure. 

17 

18 .. code-block:: python 

19 

20 @TimeTrackingWrapper 

21 def sleeper(): 

22 print("start") 

23 time.sleep(1) 

24 print("end") 

25 

26 >>> sleeper() 

27 start 

28 end 

29 INFO: foo finished after 00:00:01 (hh:mm:ss) 

30 

31 """ 

32 

33 def __init__(self, func): 

34 """Construct.""" 

35 wraps(func)(self) 

36 

37 def __call__(self, *args, **kwargs): 

38 """Start time tracking.""" 

39 with TimeTracking(name=self.__wrapped__.__name__): 

40 return self.__wrapped__(*args, **kwargs) 

41 

42 def __get__(self, instance, cls): 

43 """Create bound method object and supply self argument to the decorated method.""" 

44 if instance is None: 44 ↛ 45line 44 didn't jump to line 45, because the condition on line 44 was never true

45 return self 

46 else: 

47 return types.MethodType(self, instance) 

48 

49 

50class TimeTracking(object): 

51 """ 

52 Track time to measure execution time. 

53 

54 Time tracking automatically starts on initialisation and ends by calling stop method. Duration can always be shown 

55 by printing the time tracking object or calling get_current_duration. It is possible to start and stop time tracking 

56 by hand like 

57 

58 .. code-block:: python 

59 

60 time = TimeTracking(start=True) # start=True is default and not required to set 

61 do_something() 

62 time.stop(get_duration=True) 

63 

64 A more comfortable way is to use TimeTracking in a with statement like: 

65 

66 .. code-block:: python 

67 

68 with TimeTracking(): 

69 do_something() 

70 

71 The only disadvantage of the latter implementation is, that the duration is logged but not returned. 

72 """ 

73 

74 def __init__(self, start=True, name="undefined job", logging_level=logging.INFO, log_on_enter=False): 

75 """Construct time tracking and start if enabled.""" 

76 self.start = None 

77 self.end = None 

78 self._name = name 

79 self._logging = {logging.INFO: logging.info, logging.DEBUG: logging.debug}.get(logging_level, logging.info) 

80 self._log_on_enter = log_on_enter 

81 if start: 

82 self._start() 

83 

84 def _start(self) -> None: 

85 """Start time tracking.""" 

86 self.start = time.time() 

87 self.end = None 

88 

89 def _end(self) -> None: 

90 """Stop time tracking.""" 

91 self.end = time.time() 

92 

93 def _duration(self) -> float: 

94 """Get duration in seconds.""" 

95 if self.end: 

96 return self.end - self.start 

97 else: 

98 return time.time() - self.start 

99 

100 def __repr__(self) -> str: 

101 """Display current passed time.""" 

102 return f"{dt.timedelta(seconds=math.ceil(self._duration()))} (hh:mm:ss)" 

103 

104 def run(self) -> None: 

105 """Start time tracking.""" 

106 self._start() 

107 

108 def stop(self, get_duration=False) -> Optional[float]: 

109 """ 

110 Stop time tracking. 

111 

112 Will raise an error if time tracking was already stopped. 

113 :param get_duration: return passed time if enabled. 

114 

115 :return: duration if enabled or None 

116 """ 

117 if self.end is None: 

118 self._end() 

119 else: 

120 msg = f"Time was already stopped {time.time() - self.end}s ago." 

121 raise AssertionError(msg) 

122 if get_duration: 

123 return self.duration() 

124 

125 def duration(self) -> float: 

126 """Return duration in seconds.""" 

127 return self._duration() 

128 

129 def __enter__(self): 

130 """Context manager.""" 

131 self._logging(f"start {self._name}") if self._log_on_enter is True else None 

132 return self 

133 

134 def __exit__(self, exc_type, exc_val, exc_tb) -> None: 

135 """Stop time tracking on exit and log info about passed time.""" 

136 self.stop() 

137 self._logging(f"{self._name} finished after {self}")