muutils.logger.timing

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