codekingpro/portable-devtools

1"""2Thin wrappers around `concurrent.futures`.3"""4from contextlib import contextmanager5from operator import length_hint6from os import cpu_count7 8from ..auto import tqdm as tqdm_auto9from ..std import TqdmWarning10 11__author__ = {"github.com/": ["casperdcl"]}12__all__ = ['thread_map', 'process_map']13 14 15@contextmanager16def ensure_lock(tqdm_class, lock_name=""):17    """get (create if necessary) and then restore `tqdm_class`'s lock"""18    old_lock = getattr(tqdm_class, '_lock', None)  # don't create a new lock19    lock = old_lock or tqdm_class.get_lock()  # maybe create a new lock20    lock = getattr(lock, lock_name, lock)  # maybe subtype21    tqdm_class.set_lock(lock)22    yield lock23    if old_lock is None:24        del tqdm_class._lock25    else:26        tqdm_class.set_lock(old_lock)27 28 29def _executor_map(PoolExecutor, fn, *iterables, **tqdm_kwargs):30    """31    Implementation of `thread_map` and `process_map`.32 33    Parameters34    ----------35    tqdm_class  : [default: tqdm.auto.tqdm].36    max_workers  : [default: min(32, cpu_count() + 4)].37    chunksize  : [default: 1].38    lock_name  : [default: "":str].39    """40    kwargs = tqdm_kwargs.copy()41    if "total" not in kwargs:42        kwargs["total"] = length_hint(iterables[0])43    tqdm_class = kwargs.pop("tqdm_class", tqdm_auto)44    max_workers = kwargs.pop("max_workers", min(32, cpu_count() + 4))45    chunksize = kwargs.pop("chunksize", 1)46    lock_name = kwargs.pop("lock_name", "")47    with ensure_lock(tqdm_class, lock_name=lock_name) as lk:48        # share lock in case workers are already using `tqdm`49        with PoolExecutor(max_workers=max_workers, initializer=tqdm_class.set_lock,50                          initargs=(lk,)) as ex:51            return list(tqdm_class(ex.map(fn, *iterables, chunksize=chunksize), **kwargs))52 53 54def thread_map(fn, *iterables, **tqdm_kwargs):55    """56    Equivalent of `list(map(fn, *iterables))`57    driven by `concurrent.futures.ThreadPoolExecutor`.58 59    Parameters60    ----------61    tqdm_class  : optional62        `tqdm` class to use for bars [default: tqdm.auto.tqdm].63    max_workers  : int, optional64        Maximum number of workers to spawn; passed to65        `concurrent.futures.ThreadPoolExecutor.__init__`.66        [default: max(32, cpu_count() + 4)].67    """68    from concurrent.futures import ThreadPoolExecutor69    return _executor_map(ThreadPoolExecutor, fn, *iterables, **tqdm_kwargs)70 71 72def process_map(fn, *iterables, **tqdm_kwargs):73    """74    Equivalent of `list(map(fn, *iterables))`75    driven by `concurrent.futures.ProcessPoolExecutor`.76 77    Parameters78    ----------79    tqdm_class  : optional80        `tqdm` class to use for bars [default: tqdm.auto.tqdm].81    max_workers  : int, optional82        Maximum number of workers to spawn; passed to83        `concurrent.futures.ProcessPoolExecutor.__init__`.84        [default: min(32, cpu_count() + 4)].85    chunksize  : int, optional86        Size of chunks sent to worker processes; passed to87        `concurrent.futures.ProcessPoolExecutor.map`. [default: 1].88    lock_name  : str, optional89        Member of `tqdm_class.get_lock()` to use [default: mp_lock].90    """91    from concurrent.futures import ProcessPoolExecutor92    if iterables and "chunksize" not in tqdm_kwargs:93        # default `chunksize=1` has poor performance for large iterables94        # (most time spent dispatching items to workers).95        longest_iterable_len = max(map(length_hint, iterables))96        if longest_iterable_len > 1000:97            from warnings import warn98            warn("Iterable length %d > 1000 but `chunksize` is not set."99                 " This may seriously degrade multiprocess performance."100                 " Set `chunksize=1` or more." % longest_iterable_len,101                 TqdmWarning, stacklevel=2)102    if "lock_name" not in tqdm_kwargs:103        tqdm_kwargs = tqdm_kwargs.copy()104        tqdm_kwargs["lock_name"] = "mp_lock"105    return _executor_map(ProcessPoolExecutor, fn, *iterables, **tqdm_kwargs)106