Run a task in parallel with progress - ThreadPoolExecutor with TQDM

thread_progress.py

from typing import Callable, Iterable, Dict, List
from concurrent.futures import ThreadPoolExecutor
from tqdm import tqdm

def thread_progress(
    worker: Callable,
    items: Iterable,
    threads: int = None,
    result_handler: Callable = None,
    worker_kwargs: Dict = None,
    tqdm_kwargs: Dict = None,
) -> List:
    """
    Parallelize a task using threads and progress bar.

    :param worker: any callable (function, method, etc.)
    :param items: a list of items to run worker on.
    :param threads: number of threads to use. defaults to python version default.
    :param result_handler: any callable (function, method, etc) which will be run
        on thread results as soon as they complete.
    :param worker_kwargs: dict of keyword arguments to pass to worker.
    :param tqdm_kwargs: dict of keyword arguments to pass to tqdm (progress bar)
    :return: list of results returned by worker for all items.
    """
    results = []
    worker_kwargs = worker_kwargs or {}
    tqdm_kwargs = tqdm_kwargs or {}
    with ThreadPoolExecutor(max_workers=threads) as executor:
        futures = [executor.submit(worker, item, **worker_kwargs) for item in items]
        for future in tqdm(as_completed(futures), total=len(items), **tqdm_kwargs):
            result = future.result()
            results.append(result)
            if result_handler is not None:
                result_handler(result)
    return results