|
| 1 | +from collections.abc import Iterable, Iterator |
| 2 | +from sys import stderr |
| 3 | + |
| 4 | + |
| 5 | +def progress[T](items: Iterable[T], desc: str = "", total: int = 0) -> Iterator[T]: |
| 6 | + """ |
| 7 | + A simple progress iterator that yields items from the given iterable while |
| 8 | + displaying a progress indicator in place on a single line of stderr. The output is |
| 9 | + not written to stdout, so the output of the program remains clean (see doctests). |
| 10 | + |
| 11 | + for item in progress(range(1_000), desc="Processing"): |
| 12 | + process(item) |
| 13 | + |
| 14 | + Args: |
| 15 | + items: The iterable of items to process. |
| 16 | + desc: A description to display alongside the progress. Defaults to "". |
| 17 | + total: The total number of items, defaults to 0. If 0, it will be inferred from |
| 18 | + the iterable if possible. |
| 19 | + |
| 20 | + Yields: |
| 21 | + Iterator[T]: The items from the iterable, one by one. |
| 22 | + |
| 23 | + >>> tuple(progress(range(5))) |
| 24 | + (0, 1, 2, 3, 4) |
| 25 | + >>> tuple(progress(range(5), desc="Processing", total=3)) |
| 26 | + (0, 1, 2, 3, 4) |
| 27 | + >>> tuple(progress(range(5), desc="Processing", total=10)) |
| 28 | + (0, 1, 2, 3, 4) |
| 29 | + >>> tuple(progress(range(5), desc="Processing", total=-5)) |
| 30 | + (0, 1, 2, 3, 4) |
| 31 | + >>> from string import printable |
| 32 | + >>> tuple(progress(printable, desc="Printable")) # doctest: +ELLIPSIS |
| 33 | + ('0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'a', 'b', 'c', 'd', 'e', 'f',... |
| 34 | + """ |
| 35 | + total = max(total, 0) |
| 36 | + if not total and hasattr(items, "__len__"): |
| 37 | + total = len(items) # type: ignore[invalid-argument-type] |
| 38 | + |
| 39 | + for i, item in enumerate(items, 1): |
| 40 | + suffix = f"{i:,}/{total:,}" if total else f"{i:,}" |
| 41 | + print(f"\r\033[K{desc}: {suffix}", end="", file=stderr, flush=True) |
| 42 | + yield item |
| 43 | + |
| 44 | + print("\r\033[K", end="", file=stderr, flush=True) |
| 45 | + |
| 46 | + |
| 47 | +if __name__ == "__main__": |
| 48 | + import time |
| 49 | + |
| 50 | + print("start") |
| 51 | + for _item in progress(range(1_000), desc="Processing"): |
| 52 | + time.sleep(0.02) |
| 53 | + |
| 54 | + print("stop") |
0 commit comments