diff --git a/esutil/pbar.py b/esutil/pbar.py index 7f86faf..521fb3c 100644 --- a/esutil/pbar.py +++ b/esutil/pbar.py @@ -3,13 +3,13 @@ https://github.com/noamraph/tqdm """ -__all__ = ['PBar', 'prange'] +__all__ = ['PBar', 'pbar', 'prange'] import sys import time -def PBar(iterable, desc='', total=None, leave=True, file=sys.stderr, +def pbar(iterable, desc='', total=None, leave=True, file=sys.stderr, mininterval=0.5, miniters=1, n_bars=20): """ Get an iterable object, and return an iterator which acts exactly like the @@ -90,19 +90,60 @@ def PBar(iterable, desc='', total=None, leave=True, file=sys.stderr, file.write('\n') +PBar = pbar + + def prange(*args, **kwargs): """ - A shortcut for writing PBar(range(...)) + A shortcut for writing pbar(range(...)) + + Parameters + ---------- + Same args as for range. Extra keywords are sent to + Pbar e.g. - import time - from pbar import prange for i in prange(20): print(i) time.sleep(0.1) """ - return PBar(range(*args), **kwargs) + return pbar(range(*args), **kwargs) + + +def pmap(fn, iterable, chunksize=1, nproc=1, **kw): + """ + Execute the function on the inputs using multiple processes, while showing + a progress bar. The result is equivalent to doing + + list(map(fn, iterable)) + + Parameters + ---------- + fn: function + The function to execute + iterable: iterable data + The data over which to iterate + chunksize: int, optional + Default 1. It is often must faster to send large + chunks of data rather than 1. + nproc: int, optional + Number of processes to use, default 1 + **kw: + Additional keyword arguments for the progress bar. + See pbar for details + + Returns + ------- + An list of data, the equivalent of + list(map(fn, iterable)) + """ + from concurrent.futures import ProcessPoolExecutor + + with ProcessPoolExecutor(max_workers=nproc) as ex: + res = list(pbar(ex.map(fn, iterable, chunksize=chunksize), **kw)) + + return res def format_interval(t): @@ -131,10 +172,25 @@ def format_meter(n, total, elapsed, n_bars=20): percentage = '%3d%%' % (frac * 100) + if elapsed > 0: + it_per_second = n / elapsed # iterations per second + if it_per_second > 1: + rate_str = f'{it_per_second:.3g} it/s' + else: + second_per_it = elapsed / n + rate_str = f'{second_per_it:.3g} s/it' + else: + rate_str = '---' + left_str = format_interval(elapsed / n * (total-n)) if n else '?' - return '|%s| %d/%d %s [elapsed: %s left: %s]' % ( - bar, n, total, percentage, elapsed_str, left_str) + totstr = str(total) + nfmt = '%' + str(len(totstr)) + 'd' + meter_fmt = '|%s| ' + nfmt + '/' + nfmt + ' %s [%s<%s %s]' + + return meter_fmt % ( + bar, n, total, percentage, elapsed_str, left_str, rate_str + ) else: return '%d [elapsed: %s]' % (n, elapsed_str)