From 32cde4d16a445fb523ec695abc7a5f1ad2ea32ff Mon Sep 17 00:00:00 2001 From: Casper da Costa-Luis Date: Fri, 24 Jan 2020 21:43:43 +0000 Subject: [PATCH] contrib:concurrent:add initargs --- tqdm/contrib/concurrent.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tqdm/contrib/concurrent.py b/tqdm/contrib/concurrent.py index e5219681..9eb10db8 100644 --- a/tqdm/contrib/concurrent.py +++ b/tqdm/contrib/concurrent.py @@ -12,6 +12,7 @@ except ImportError: except ImportError: def cpu_count(): return 4 +import sys __author__ = {"github.com/": ["casperdcl"]} __all__ = ['thread_map', 'process_map'] @@ -28,7 +29,12 @@ def _executor_map(PoolExecutor, fn, *iterables, **tqdm_kwargs): kwargs.setdefault("total", len(iterables[0])) tqdm_class = kwargs.pop("tqdm_class", tqdm_auto) max_workers = kwargs.pop("max_workers", min(32, cpu_count() + 4)) - with PoolExecutor(max_workers=max_workers) as ex: + pool_kwargs = dict(max_workers=max_workers) + if sys.version_info[:2] >= (3, 7): + # share lock in case workers are already using `tqdm` + pool_kwargs.update( + initializer=tqdm_class.set_lock, initargs=(tqdm_class.get_lock(),)) + with PoolExecutor(**pool_kwargs) as ex: return list(tqdm_class(ex.map(fn, *iterables), **kwargs))