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

1from __future__ import annotations 

2 

3import time 

4from typing import Literal 

5 

6 

7class TimerContext: 

8 """context manager for timing code""" 

9 

10 def __init__(self) -> None: 

11 self.start_time: float 

12 self.end_time: float 

13 self.elapsed_time: float 

14 

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

16 self.start_time = time.time() 

17 return self 

18 

19 def __exit__(self, exc_type, exc_val, exc_tb) -> Literal[False]: 

20 self.end_time = time.time() 

21 self.elapsed_time = self.end_time - self.start_time 

22 return False 

23 

24 

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

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

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

28 return time[3:] 

29 else: 

30 return time 

31 

32 

33class ProgressEstimator: 

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

35 

36 def __init__( 

37 self, 

38 n_total: int, 

39 pbar_fill: str = "█", 

40 pbar_empty: str = " ", 

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

42 ): 

43 self.n_total: int = n_total 

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

45 self.pbar_fill: str = pbar_fill 

46 self.pbar_empty: str = pbar_empty 

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

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

49 

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

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

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

53 per_iter: float = elapsed / i 

54 return dict( 

55 elapsed=elapsed, 

56 per_iter=per_iter, 

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

58 percent=i / self.n_total, 

59 ) 

60 

61 def get_pbar( 

62 self, 

63 i: int, 

64 width: int = 30, 

65 ) -> str: 

66 """returns a progress bar""" 

67 percent_filled: float = i / self.n_total 

68 # round to nearest integer 

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

70 return "".join( 

71 [ 

72 self.pbar_bounds[0], 

73 self.pbar_fill * n_filled, 

74 self.pbar_empty * (width - n_filled), 

75 self.pbar_bounds[1], 

76 ] 

77 ) 

78 

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

80 """returns a progress string""" 

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

82 

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

84 # TODO: get_progress_default 

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

86 # timing_str: str 

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