Simplification of parallelisation.

This commit is contained in:
Nick Brassel 2023-09-29 11:25:53 +10:00
parent 2808171cad
commit 06316a1cd4
No known key found for this signature in database

View File

@ -25,6 +25,24 @@ def _set_log_level(level):
return old return old
@contextlib.contextmanager
def parallelize(parallel):
with contextlib.suppress(ImportError):
from mpire import WorkerPool
if parallel:
with WorkerPool() as pool:
yield functools.partial(pool.imap_unordered, progress_bar=True)
else:
yield map
return
if parallel:
with multiprocessing.Pool() as pool:
yield pool.imap_unordered
else:
yield map
@contextlib.contextmanager @contextlib.contextmanager
def ignore_logging(): def ignore_logging():
old = _set_log_level(logging.CRITICAL) old = _set_log_level(logging.CRITICAL)
@ -47,11 +65,15 @@ def _keymap_exists(keyboard, keymap):
return keyboard if qmk.keymap.locate_keymap(keyboard, keymap) is not None else None return keyboard if qmk.keymap.locate_keymap(keyboard, keymap) is not None else None
def _load_keymap_info(kb_km): def _load_keymap_info(arg0, arg1=None):
"""Returns a tuple of (keyboard, keymap, info.json) for the given keyboard/keymap combination. """Returns a tuple of (keyboard, keymap, info.json) for the given keyboard/keymap combination.
Caters for the different unpacking requirements of each variatn of imap_unordered().
""" """
with ignore_logging(): with ignore_logging():
return (kb_km[0], kb_km[1], keymap_json(kb_km[0], kb_km[1])) if arg1 is None:
return (arg0[0], arg0[1], keymap_json(arg0[0], arg0[1]))
return (arg0, arg1, keymap_json(arg0, arg1))
def expand_make_targets(targets: List[str], parallel=True) -> List[Tuple[str, str]]: def expand_make_targets(targets: List[str], parallel=True) -> List[Tuple[str, str]]:
@ -78,25 +100,18 @@ def _expand_keymap_target(keyboard: str, keymap: str, all_keyboards: List[str] =
all_keyboards = qmk.keyboard.list_keyboards() all_keyboards = qmk.keyboard.list_keyboards()
if keyboard == 'all': if keyboard == 'all':
if keymap == 'all':
def _inner_func(pool): cli.log.info('Retrieving list of all keyboards and keymaps...')
_map_func = pool.imap_unordered if pool is not None else map targets = []
if keymap == 'all': with parallelize(parallel) as map_func:
cli.log.info('Retrieving list of all keyboards and keymaps...') for kb in map_func(_all_keymaps, all_keyboards):
targets = []
for kb in _map_func(_all_keymaps, all_keyboards):
targets.extend(kb) targets.extend(kb)
return targets return targets
else: else:
cli.log.info(f'Retrieving list of keyboards with keymap "{keymap}"...') cli.log.info(f'Retrieving list of keyboards with keymap "{keymap}"...')
keyboard_filter = functools.partial(_keymap_exists, keymap=keymap) keyboard_filter = functools.partial(_keymap_exists, keymap=keymap)
return [(kb, keymap) for kb in filter(lambda e: e is not None, _map_func(keyboard_filter, all_keyboards))] with parallelize(parallel) as map_func:
return [(kb, keymap) for kb in filter(lambda e: e is not None, map_func(keyboard_filter, all_keyboards))]
if parallel:
with multiprocessing.Pool() as pool:
return _inner_func(pool)
return _inner_func(None)
else: else:
if keymap == 'all': if keymap == 'all':
keyboard = qmk.keyboard.resolve_keyboard(keyboard) keyboard = qmk.keyboard.resolve_keyboard(keyboard)
@ -126,15 +141,8 @@ def _filter_keymap_targets(target_list: List[Tuple[str, str]], filters: List[str
else: else:
cli.log.info('Parsing data for all matching keyboard/keymap combinations...') cli.log.info('Parsing data for all matching keyboard/keymap combinations...')
def _inner_func(pool): with parallelize(parallel) as map_func:
_map_func = pool.imap_unordered if pool is not None else map valid_keymaps = [(e[0], e[1], dotty(e[2])) for e in map_func(_load_keymap_info, target_list)]
return [(e[0], e[1], dotty(e[2])) for e in _map_func(_load_keymap_info, target_list)]
if parallel:
with multiprocessing.Pool() as pool:
valid_keymaps = _inner_func(pool)
else:
valid_keymaps = _inner_func(None)
function_re = re.compile(r'^(?P<function>[a-zA-Z]+)\((?P<key>[a-zA-Z0-9_\.]+)(,\s*(?P<value>[^#]+))?\)$') function_re = re.compile(r'^(?P<function>[a-zA-Z]+)\((?P<key>[a-zA-Z0-9_\.]+)(,\s*(?P<value>[^#]+))?\)$')
equals_re = re.compile(r'^(?P<key>[a-zA-Z0-9_\.]+)\s*=\s*(?P<value>[^#]+)$') equals_re = re.compile(r'^(?P<key>[a-zA-Z0-9_\.]+)\s*=\s*(?P<value>[^#]+)$')