Coverage for muutils / logger / timing.py: 55%

1from __future__ import annotations 

2 

3import time 

4from types import TracebackType 

5from typing import Literal 

6 

7 

8class TimerContext: 

9 """context manager for timing code""" 

10 

11 def __init__(self) -> None: 

12 self.start_time: float 

13 self.end_time: float 

14 self.elapsed_time: float 

15 

16 def __enter__(self) -> "TimerContext": 

17 self.start_time = time.time() 

18 return self 

19 

20 def __exit__( 

21 self, 

22 exc_type: type[BaseException] | None, 

23 exc_val: BaseException | None, 

24 exc_tb: TracebackType | None, 

25 ) -> Literal[False]: 

26 self.end_time = time.time() 

27 self.elapsed_time = self.end_time - self.start_time 

28 return False 

29 

30 

31def filter_time_str(time: str) -> str: 

32 """assuming format `h:mm:ss`, clips off the hours if its 0""" 

33 if (len(time) == 7) and (time[0] == "0"): 

34 return time[3:] 

35 else: 

36 return time 

37 

38 

39class ProgressEstimator: 

40 """estimates progress and can give a progress bar""" 

41 

42 def __init__( 

43 self, 

44 n_total: int, 

45 pbar_fill: str = "█", 

46 pbar_empty: str = " ", 

47 pbar_bounds: tuple[str, str] = ("|", "|"), 

48 ): 

49 self.n_total: int = n_total 

50 self.starttime: float = time.time() 

51 self.pbar_fill: str = pbar_fill 

52 self.pbar_empty: str = pbar_empty 

53 self.pbar_bounds: tuple[str, str] = pbar_bounds 

54 self.total_str_len: int = len(str(n_total)) 

55 

56 def get_timing_raw(self, i: int) -> dict[str, float]: 

57 """returns dict(elapsed, per_iter, remaining, percent)""" 

58 elapsed: float = time.time() - self.starttime 

59 per_iter: float = elapsed / i 

60 return dict( 

61 elapsed=elapsed, 

62 per_iter=per_iter, 

63 remaining=(self.n_total - i) * per_iter, 

64 percent=i / self.n_total, 

65 ) 

66 

67 def get_pbar( 

68 self, 

69 i: int, 

70 width: int = 30, 

71 ) -> str: 

72 """returns a progress bar""" 

73 percent_filled: float = i / self.n_total 

74 # round to nearest integer 

75 n_filled: int = int(round(percent_filled * width)) 

76 return "".join( 

77 [ 

78 self.pbar_bounds[0], 

79 self.pbar_fill * n_filled, 

80 self.pbar_empty * (width - n_filled), 

81 self.pbar_bounds[1], 

82 ] 

83 ) 

84 

85 def get_progress_default(self, i: int) -> str: 

86 """returns a progress string""" 

87 timing_raw: dict[str, float] = self.get_timing_raw(i) 

88 

89 percent_str: str = str(int(timing_raw["percent"] * 100)).ljust(2) 

90 # TODO: get_progress_default 

91 # iters_str: str = f"{str(i).ljust(self.total_str_len)}/{self.n_total}" 

92 # timing_str: str 

93 return f"{percent_str}% {self.get_pbar(i)}"