Move parallelize to util.py.

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)

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