1 from __future__ import annotations

2

3 import time

4 from typing import Literal

5

6

7 class 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

25 def 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

33 class 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)}"