printbuddies.printbuddies

  1from os import get_terminal_size
  2from time import sleep
  3from typing import Any
  4
  5from noiftimer import Timer
  6
  7
  8def clear():
  9    """Erase the current line from the terminal."""
 10    try:
 11        print(" " * (get_terminal_size().columns - 1), flush=True, end="\r")
 12    except OSError:
 13        ...
 14    except Exception as e:
 15        raise e
 16
 17
 18def print_in_place(string: str, animate: bool = False, animate_refresh: float = 0.01):
 19    """Calls to `print_in_place` will overwrite the previous line of text in the terminal with `string`.
 20
 21    #### :params:
 22
 23    `animate`: Will cause `string` to be printed to the terminal one character at a time.
 24
 25    `animate_refresh`: Number of seconds between the addition of characters when `animate` is `True`."""
 26    clear()
 27    string = str(string)
 28    try:
 29        width = get_terminal_size().columns
 30        string = string[: width - 2]
 31        if animate:
 32            for i in range(len(string)):
 33                print(f"{string[:i+1]}", flush=True, end=" \r")
 34                sleep(animate_refresh)
 35        else:
 36            print(string, flush=True, end="\r")
 37    except OSError:
 38        ...
 39    except Exception as e:
 40        raise e
 41
 42
 43def ticker(info: list[str]):
 44    """Prints `info` to terminal with top and bottom padding so that previous text is not visible.
 45
 46    Similar visually to `print_in_place`, but for multiple lines."""
 47    try:
 48        width = get_terminal_size().columns
 49        info = [str(line)[: width - 1] for line in info]
 50        height = get_terminal_size().lines - len(info)
 51        print("\n" * (height * 2), end="")
 52        print(*info, sep="\n", end="")
 53        print("\n" * (int((height) / 2)), end="")
 54    except OSError:
 55        ...
 56    except Exception as e:
 57        raise e
 58
 59
 60class ProgBar:
 61    """Self incrementing, dynamically sized progress bar.
 62
 63    Includes an internal timer that starts when this object is created.
 64
 65    Easily add runtime to progress display:
 66
 67    >>> bar = ProgBar(total=100)
 68    >>> time.sleep(30)
 69    >>> bar.display(prefix=f"Doin stuff ~ {bar.runtime}")
 70    >>> "Doin stuff ~ runtime: 30s [_///////////////////]-1.00%" """
 71
 72    def __init__(
 73        self,
 74        total: float,
 75        update_frequency: int = 1,
 76        fill_ch: str = "_",
 77        unfill_ch: str = "/",
 78        width_ratio: float = 0.5,
 79        new_line_after_completion: bool = True,
 80        clear_after_completion: bool = False,
 81    ):
 82        """
 83        #### :params:
 84
 85        `total`: The number of calls to reach 100% completion.
 86
 87        `update_frequency`: The progress bar will only update once every this number of calls to `display()`.
 88        The larger the value, the less performance impact `ProgBar` has on the loop in which it is called.
 89        e.g.
 90        >>> bar = ProgBar(100, update_frequency=10)
 91        >>> for _ in range(100):
 92        >>>     bar.display()
 93
 94        ^The progress bar in the terminal will only update once every ten calls, going from 0%->100% in 10% increments.
 95        Note: If `total` is not a multiple of `update_frequency`, the display will not show 100% completion when the loop finishes.
 96
 97        `fill_ch`: The character used to represent the completed part of the bar.
 98
 99        `unfill_ch`: The character used to represent the incomplete part of the bar.
100
101        `width_ratio`: The width of the progress bar relative to the width of the terminal window.
102
103        `new_line_after_completion`: Make a call to `print()` once `self.counter >= self.total`.
104
105        `clear_after_completion`: Make a call to `printbuddies.clear()` once `self.counter >= self.total`.
106
107        Note: if `new_line_after_completion` and `clear_after_completion` are both `True`, the line will be cleared
108        then a call to `print()` will be made."""
109        self.total = total
110        self.update_frequency = update_frequency
111        self.fill_ch = fill_ch[0]
112        self.unfill_ch = unfill_ch[0]
113        self.width_ratio = width_ratio
114        self.new_line_after_completion = new_line_after_completion
115        self.clear_after_completion = clear_after_completion
116        self.reset()
117        self.with_context = False
118
119    def __enter__(self):
120        self.with_context = True
121        return self
122
123    def __exit__(self, *args, **kwargs):
124        if self.clear_after_completion:
125            clear()
126        else:
127            print()
128
129    def reset(self):
130        self.counter = 1
131        self.percent = ""
132        self.prefix = ""
133        self.suffix = ""
134        self.filled = ""
135        self.unfilled = ""
136        self.timer = Timer(subsecond_resolution=False).start()
137
138    @property
139    def runtime(self) -> str:
140        return f"runtime:{self.timer.elapsed_str}"
141
142    @property
143    def bar(self) -> str:
144        return f"{self.prefix}{' '*bool(self.prefix)}[{self.filled}{self.unfilled}]-{self.percent}% {self.suffix}"
145
146    def get_percent(self) -> str:
147        """Returns the percentage completed to two decimal places as a string without the `%`."""
148        percent = str(round(100.0 * self.counter / self.total, 2))
149        if len(percent.split(".")[1]) == 1:
150            percent = percent + "0"
151        if len(percent.split(".")[0]) == 1:
152            percent = "0" + percent
153        return percent
154
155    def _prepare_bar(self):
156        self.terminal_width = get_terminal_size().columns - 1
157        bar_length = int(self.terminal_width * self.width_ratio)
158        progress = int(bar_length * min(self.counter / self.total, 1.0))
159        self.filled = self.fill_ch * progress
160        self.unfilled = self.unfill_ch * (bar_length - progress)
161        self.percent = self.get_percent()
162
163    def _trim_bar(self):
164        original_width = self.width_ratio
165        while len(self.bar) > self.terminal_width and self.width_ratio > 0:
166            self.width_ratio -= 0.01
167            self._prepare_bar()
168        self.width_ratio = original_width
169
170    def get_bar(self):
171        return f"{self.prefix}{' '*bool(self.prefix)}[{self.filled}{self.unfilled}]-{self.percent}% {self.suffix}"
172
173    def display(
174        self,
175        prefix: str = "",
176        suffix: str = "",
177        counter_override: float | None = None,
178        total_override: float | None = None,
179        return_object: Any | None = None,
180    ) -> Any:
181        """Writes the progress bar to the terminal.
182
183        #### :params:
184
185        `prefix`: String affixed to the front of the progress bar.
186
187        `suffix`: String appended to the end of the progress bar.
188
189        `counter_override`: When an externally incremented completion counter is needed.
190
191        `total_override`: When an externally controlled bar total is needed.
192
193        `return_object`: An object to be returned by display().
194        Allows `display()` to be called within a comprehension:
195
196        e.g.
197
198        >>> bar = ProgBar(10)
199        >>> def square(x: int | float)->int|float:
200        >>>     return x * x
201        >>> myList = [bar.display(return_object=square(i)) for i in range(10)]
202        >>> <progress bar gets displayed>
203        >>> myList
204        >>> [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]"""
205        if not self.timer.started:
206            self.timer.start()
207        if counter_override is not None:
208            self.counter = counter_override
209        if total_override:
210            self.total = total_override
211        # Don't wanna divide by 0 there, pal
212        while self.total <= 0:
213            self.total += 1
214        try:
215            if self.counter % self.update_frequency == 0:
216                self.prefix = prefix
217                self.suffix = suffix
218                self._prepare_bar()
219                self._trim_bar()
220                pad = " " * (self.terminal_width - len(self.bar))
221                width = get_terminal_size().columns
222                print(f"{self.bar}{pad}"[: width - 2], flush=True, end="\r")
223            if self.counter >= self.total:
224                self.timer.stop()
225                if not self.with_context:
226                    if self.clear_after_completion:
227                        clear()
228                    if self.new_line_after_completion:
229                        print()
230            self.counter += 1
231        except OSError:
232            ...
233        except Exception as e:
234            raise e
235        return return_object
236
237
238class Spinner:
239    """Prints one of a sequence of characters in order everytime `display()` is called.
240
241    The `display` function writes the new character to the same line, overwriting the previous character.
242
243    The sequence will be cycled through indefinitely.
244
245    If used as a context manager, the last printed character will be cleared upon exiting.
246    """
247
248    def __init__(
249        self, sequence: list[str] = ["/", "-", "\\"], width_ratio: float = 0.25
250    ):
251        """
252        #### params:
253
254        `sequence`: Override the built in spin sequence.
255
256        `width_ratio`: The fractional amount of the terminal for characters to move across."""
257        self._base_sequence = sequence
258        self.width_ratio = width_ratio
259        self.sequence = self._base_sequence
260
261    def __enter__(self):
262        return self
263
264    def __exit__(self, *args, **kwargs):
265        clear()
266
267    @property
268    def width_ratio(self) -> float:
269        return self._width_ratio
270
271    @width_ratio.setter
272    def width_ratio(self, ratio: float):
273        self._width_ratio = ratio
274        self._update_width()
275
276    def _update_width(self):
277        self._current_terminal_width = get_terminal_size().columns
278        self._width = int((self._current_terminal_width - 1) * self.width_ratio)
279
280    @property
281    def sequence(self) -> list[Any]:
282        return self._sequence
283
284    @sequence.setter
285    def sequence(self, character_list: list[Any]):
286        self._sequence = [
287            ch.rjust(i + j)
288            for i in range(1, self._width, len(character_list))
289            for j, ch in enumerate(character_list)
290        ]
291        self._sequence += self._sequence[::-1]
292
293    def _get_next(self) -> str:
294        """Pop the first element of `self._sequence`, append it to the end, and return the element."""
295        ch = self.sequence.pop(0)
296        self.sequence.append(ch)
297        return ch
298
299    def display(self):
300        """Print the next character in the sequence."""
301        try:
302            if get_terminal_size().columns != self._current_terminal_width:
303                self._update_width()
304                self.sequence = self._base_sequence
305            print_in_place(self._get_next())
306        except OSError:
307            ...
308        except Exception as e:
309            raise e