diff --git a/lib/python/qmk/search.py b/lib/python/qmk/search.py index 4c61f332881..9acc59668f6 100644 --- a/lib/python/qmk/search.py +++ b/lib/python/qmk/search.py @@ -4,12 +4,12 @@ import contextlib import functools import fnmatch import logging -import multiprocessing import re from typing import List, Tuple from dotty_dict import dotty from milc import cli +from qmk.util import parallelize from qmk.info import keymap_json import qmk.keyboard import qmk.keymap @@ -25,23 +25,6 @@ def _set_log_level(level): return old -@contextlib.contextmanager -def parallelize(parallel): - if not parallel: - yield map - return - - with contextlib.suppress(ImportError): - from mpire import WorkerPool - with WorkerPool() as pool: - yield functools.partial(pool.imap_unordered, progress_bar=True) - return - - with multiprocessing.Pool() as pool: - yield pool.imap_unordered - return - - @contextlib.contextmanager def ignore_logging(): old = _set_log_level(logging.CRITICAL) diff --git a/lib/python/qmk/util.py b/lib/python/qmk/util.py new file mode 100644 index 00000000000..d57bce951c9 --- /dev/null +++ b/lib/python/qmk/util.py @@ -0,0 +1,27 @@ +"""Utility functions. +""" +import contextlib +import functools +import multiprocessing + +@contextlib.contextmanager +def parallelize(do_parallel): + """Returns a function that can be used in place of a map() call. + + Attempts to use `mpire`, falling back to `multiprocessing` if it's not + available. If parallelization is not requested, returns the original map() + function. + """ + if not do_parallel: + yield map + return + + with contextlib.suppress(ImportError): + from mpire import WorkerPool + with WorkerPool() as pool: + yield functools.partial(pool.imap_unordered, progress_bar=True) + return + + with multiprocessing.Pool() as pool: + yield pool.imap_unordered + return