diff options
Diffstat (limited to 'sec_certs/parallel_processing.py')
| -rw-r--r-- | sec_certs/parallel_processing.py | 16 |
1 files changed, 9 insertions, 7 deletions
diff --git a/sec_certs/parallel_processing.py b/sec_certs/parallel_processing.py index 912ef73b..c42c9c16 100644 --- a/sec_certs/parallel_processing.py +++ b/sec_certs/parallel_processing.py @@ -1,25 +1,27 @@ from sec_certs.helpers import tqdm -from concurrent.futures import ProcessPoolExecutor as ProcessPool, ThreadPoolExecutor as ThreadPool +from multiprocessing.pool import ThreadPool +from billiard.pool import Pool from typing import Callable, Iterable, Optional, Union import time -def process_parallel(func: Callable, items: Iterable, max_workers: int, +def process_parallel(func: Callable, items: Iterable, max_workers: int, callback: Optional[Callable] = None, use_threading: bool = True, progress_bar: bool = True, unpack: bool = False, progress_bar_desc: Optional[str] = None): - pool: Union[ProcessPool, ThreadPool] = ThreadPool(max_workers) if use_threading else ProcessPool(max_workers) - results = [pool.submit(func, *i) for i in items] if unpack else [pool.submit(func, i) for i in items] + pool: Union[Pool, ThreadPool] = ThreadPool(max_workers) if use_threading else Pool(max_workers) + results = [pool.apply_async(func, (*i,), callback=callback) for i in items] if unpack else [pool.apply_async(func, (i, ), callback=callback) for i in items] if progress_bar is True and items: bar = tqdm(total=len(results), desc=progress_bar_desc) - while not all(all_done := [x.done() for x in results]): + while not all(all_done := [x.ready() for x in results]): done_count = len(list(filter(lambda x: x, all_done))) bar.update(done_count - bar.n) time.sleep(1) bar.update(len(results) - bar.n) bar.close() - pool.shutdown() + pool.close() + pool.join() - return [r.result() for r in results] + return [r.get() for r in results] |
