swansonk14 / p_tqdm

Parallel processing with progress bars
MIT License
457 stars 44 forks source link

Add keyword to disable progress bar #25

Closed ChristianMichelsen closed 4 years ago

ChristianMichelsen commented 4 years ago

Hi and thanks for a great package!

I was wondering if it would make sense to add a keyword to disable the progress bar? I know that the progress bar is the reason why one would use this package to begin with, but I still sometimes run into cases where it would be useful to be able to disable the progress bar. I was thinking something along the line of:

def _parallel(ordered, function, *iterables, **kwargs):
    """Returns a generator for a parallel map with a progress bar.
    Arguments:
        ordered(bool): True for an ordered map, false for an unordered map.
        function(Callable): The function to apply to each element of the given Iterables.
        iterables(Tuple[Iterable]): One or more Iterables containing the data to be mapped.
    Returns:
        A generator which will apply the function to each element of the given Iterables
        in parallel in order with a progress bar.
    """

    # Extract num_cpus
    num_cpus = kwargs.pop('num_cpus', None)
    do_tqdm = kwargs.pop('do_tqdm', True)

    # Determine num_cpus
    if num_cpus is None:
        num_cpus = cpu_count()
    elif type(num_cpus) == float:
        num_cpus = int(round(num_cpus * cpu_count()))

    # Determine length of tqdm (equal to length of shortest iterable)
    length = min(len(iterable) for iterable in iterables if isinstance(iterable, Sized))

    # Create parallel generator
    map_type = 'imap' if ordered else 'uimap'
    pool = Pool(num_cpus)
    map_func = getattr(pool, map_type)

    # create iterable
    items = map_func(function, *iterables)

    # add progress bar
    if do_tqdm:
        items = tqdm(items, total=length, **kwargs)

    for item in items:
        yield item

    pool.clear()

How would people think about this?

Cheers, Christian

Irazall commented 4 years ago

Sounds nice!

lsdch commented 4 years ago

p_tqdm functions accept keyword args to be forwarded to tqdm (see https://github.com/tqdm/tqdm#parameters). Setting disable=True already seems to do the job :

p_map(lambda x: x^2, range(1000), disable=True) # no progress bar 
ChristianMichelsen commented 4 years ago

Oh, that's a lot smarter and easier! Why didn't I think of that before, haha? Oh well, thanks a lot for the solution!