Support parallelizing commands in an alias.

jsirois · jsirois · commit 134506c7c046 · 2024-12-24T02:54:53.000-08:00
diff --git a/CHANGES.md b/CHANGES.md
@@ -1,5 +1,9 @@
 # Release Notes
 
+## 0.3.0
+
+Add support for parallelizing execution of commands in an alias.
+
 ## 0.2.1
 
 Fix project root dir detection.
diff --git a/dev_cmd/__init__.py b/dev_cmd/__init__.py
@@ -1,4 +1,4 @@
 # Copyright 2024 John Sirois.
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
-__version__ = "0.2.1"
+__version__ = "0.3.0"
diff --git a/dev_cmd/errors.py b/dev_cmd/errors.py
@@ -16,3 +16,7 @@ class InvalidArgumentError(DevError):
 
 class InvalidModelError(DevError):
     """Indicates invalid dev command configuration."""
+
+
+class ParallelExecutionError(Exception):
+    """Conveys details of 2 or more failed parallel commands."""
diff --git a/dev_cmd/model.py b/dev_cmd/model.py
@@ -22,30 +22,47 @@ class Command:
 @dataclass(frozen=True)
 class Dev:
     commands: Mapping[str, Command]
-    aliases: Mapping[str, tuple[Command, ...]]
-    default: tuple[str, tuple[Command, ...]] | None = None
+    aliases: Mapping[str, tuple[Command | tuple[Command, ...], ...]]
+    default: tuple[str, tuple[Command | tuple[Command, ...], ...]] | None = None
     source: Any = "<code>"
 
 
+@dataclass
+class ExtraArgsChecker:
+    accepts_extra_args: Command | None = None
+
+    def check(self, *commands: Command) -> None:
+        for command in commands:
+            if command.accepts_extra_args:
+                if self.accepts_extra_args is not None:
+                    raise InvalidModelError(
+                        f"The command {command.name!r} accepts extra args, but only one "
+                        f"command can accept extra args per invocation and command "
+                        f"{self.accepts_extra_args.name!r} already does."
+                    )
+                self.accepts_extra_args = command
+
+
 @dataclass(frozen=True)
 class Invocation:
     @classmethod
-    def create(cls, *tasks: tuple[str, Iterable[Command]]) -> Invocation:
-        _tasks: dict[str, tuple[Command, ...]] = {}
-        accepts_extra_args: Command | None = None
+    def create(cls, *tasks: tuple[str, Iterable[Command | Iterable[Command]]]) -> Invocation:
+        _tasks: dict[str, tuple[Command | tuple[Command, ...], ...]] = {}
+        extra_args_checker = ExtraArgsChecker()
         for task, commands in tasks:
-            _tasks[task] = tuple(commands)
+            task_cmds: list[Command | tuple[Command, ...]] = []
             for command in commands:
-                if command.accepts_extra_args:
-                    if accepts_extra_args is not None:
-                        raise InvalidModelError(
-                            f"The command {command.name!r} accepts extra args, but only one "
-                            f"command can accept extra args per invocation and command "
-                            f"{accepts_extra_args.name!r} already does."
-                        )
-                    accepts_extra_args = command
-
-        return cls(tasks=_tasks, accepts_extra_args=accepts_extra_args is not None)
-
-    tasks: Mapping[str, tuple[Command, ...]]
+                if isinstance(command, Command):
+                    extra_args_checker.check(command)
+                    task_cmds.append(command)
+                else:
+                    extra_args_checker.check(*command)
+                    task_cmds.append(tuple(command))
+            _tasks[task] = tuple(task_cmds)
+
+        return cls(
+            tasks=_tasks, accepts_extra_args=extra_args_checker.accepts_extra_args is not None
+        )
+
+    tasks: Mapping[str, tuple[Command | tuple[Command, ...], ...]]
     accepts_extra_args: bool
diff --git a/dev_cmd/parse.py b/dev_cmd/parse.py
@@ -92,24 +92,51 @@ def _parse_commands(commands: dict[str, Any] | None, project_dir: Path) -> Itera
         yield Command(name, env, args, cwd, accepts_extra_args=accepts_extra_args)
 
 
-def _parse_aliases(aliases: dict[str, Any] | None) -> Iterator[tuple[str, tuple[str, ...]]]:
-    if aliases:
-        for alias, commands in aliases.items():
-            yield alias, tuple(_assert_list_str(commands, path=f"[tool.dev-cmd.aliases.{alias}]"))
+def _parse_aliases(
+    aliases: dict[str, Any] | None,
+) -> Iterator[tuple[str, tuple[str | tuple[str, ...], ...]]]:
+    if not aliases:
+        return
+
+    def iter_commands(alias: str, obj: Any) -> Iterator[str | tuple[str, ...]]:
+        if not isinstance(commands, list):
+            raise InvalidModelError(
+                f"Expected value at [tool.dev-cmd.aliases] `{alias}` to be a list containing "
+                f"strings or lists of strings, but given: {obj} of type {type(obj)}."
+            )
+
+        for index, item in enumerate(obj):
+            if isinstance(item, str):
+                yield item
+            elif isinstance(item, list):
+                if not all(isinstance(element, str) for element in item):
+                    raise InvalidModelError(
+                        f"Expected value at [tool.dev-cmd.aliases] `{alias}`[{index}] to be a list "
+                        f"of strings, but given list with at least one non-string item: {item}."
+                    )
+                yield tuple(item)
+            else:
+                raise InvalidModelError(
+                    f"Expected value at [tool.dev-cmd.aliases] `{alias}`[{index}] to be a string "
+                    f"or a list of strings, but given: {item} of type {type(item)}."
+                )
+
+    for alias, commands in aliases.items():
+        yield alias, tuple(iter_commands(alias, commands))
 
 
 def _parse_default(
     default: dict[str, Any] | None,
     commands: Mapping[str, Command],
-    aliases: Mapping[str, tuple[Command, ...]],
-) -> tuple[str, tuple[Command, ...]] | None:
+    aliases: Mapping[str, tuple[Command | tuple[Command, ...], ...]],
+) -> tuple[str, tuple[Command | tuple[Command, ...], ...]] | None:
     if not default:
         if len(commands) == 1:
             name, command = next(iter(commands.items()))
             return name, tuple([command])
         return None
 
-    default_commands: tuple[str, tuple[Command, ...]] | None = None
+    default_commands: tuple[str, tuple[Command | tuple[Command, ...], ...]] | None = None
     alias = default.pop("alias", None)
     if alias:
         if not isinstance(alias, str):
@@ -162,23 +189,41 @@ def pop_dict(key: str, *, path: str) -> dict[str, Any] | None:
             project_dir=pyproject_toml.path.parent,
         )
     }
-    aliases: dict[str, tuple[Command, ...]] = {}
-    for alias, cmds in _parse_aliases(pop_dict("aliases", path="[tool.dev-cmd.aliases-]")):
+    aliases: dict[str, tuple[Command | tuple[Command, ...], ...]] = {}
+    for alias, cmds in _parse_aliases(pop_dict("aliases", path="[tool.dev-cmd.aliases]")):
         if alias in commands:
             raise InvalidModelError(
                 f"The alias name {alias!r} conflicts with a command of the same name."
             )
-        alias_cmds: list[Command] = []
-        for cmd in cmds:
-            if cmd in commands:
-                alias_cmds.append(commands[cmd])
-            elif cmd in aliases:
-                alias_cmds.extend(aliases[cmd])
+        alias_cmds: list[Command | tuple[Command, ...]] = []
+        for index, cmd in enumerate(cmds):
+            if isinstance(cmd, str):
+                if cmd in commands:
+                    alias_cmds.append(commands[cmd])
+                elif cmd in aliases:
+                    alias_cmds.extend(aliases[cmd])
+                else:
+                    raise InvalidModelError(
+                        f"The task {cmd!r} defined in alias {alias!r} is neither a command nor a "
+                        f"previously defined alias."
+                    )
             else:
-                raise InvalidModelError(
-                    f"The task {cmd!r} defined in alias {alias!r} is neither a command nor a "
-                    f"previously defined alias."
-                )
+                parallel_cmds: list[Command] = []
+                for parallel_cmd in cmd:
+                    if parallel_cmd in commands:
+                        parallel_cmds.append(commands[parallel_cmd])
+                    elif parallel_cmd in aliases:
+                        raise InvalidModelError(
+                            f"Expected value at [tool.dev-cmd.aliases] `{alias}`[{index}] to be a "
+                            f"list of command names, but {parallel_cmd!r} is an alias."
+                        )
+                    else:
+                        raise InvalidModelError(
+                            f"Expected value at [tool.dev-cmd.aliases] `{alias}`[{index}] to be a "
+                            f"list of command names, but {parallel_cmd!r} is doesn't correspond "
+                            f"with any defined command."
+                        )
+                alias_cmds.append(tuple(parallel_cmds))
         aliases[alias] = tuple(alias_cmds)
 
     default = _parse_default(pop_dict("default", path="[tool.dev-cmd.default]"), commands, aliases)
diff --git a/dev_cmd/run.py b/dev_cmd/run.py
@@ -3,22 +3,91 @@
 
 from __future__ import annotations
 
+import asyncio
 import os
-import subprocess
 import sys
 from argparse import ArgumentParser
+from asyncio.subprocess import Process
 from subprocess import CalledProcessError
 from typing import Any, Iterable
 
+import aioconsole
 from colors import colors
 
 from dev_cmd import __version__
-from dev_cmd.errors import DevError, InvalidArgumentError, InvalidModelError
-from dev_cmd.model import Dev, Invocation
+from dev_cmd.errors import DevError, InvalidArgumentError, InvalidModelError, ParallelExecutionError
+from dev_cmd.model import Command, Dev, Invocation
 from dev_cmd.parse import parse_dev_config
 from dev_cmd.project import find_pyproject_toml
 
 
+async def _invoke_command(command, extra_args, **subprocess_kwargs: Any) -> Process:
+    args = list(command.args)
+    if extra_args and command.accepts_extra_args:
+        args.extend(extra_args)
+    if not os.path.exists(command.cwd):
+        raise InvalidModelError(
+            f"The `cwd` for command {command.name!r} does not exist: {command.cwd}"
+        )
+    return await asyncio.create_subprocess_exec(
+        command.args[0], *command.args[1:], **subprocess_kwargs
+    )
+
+
+async def _invoke(invocation: Invocation, extra_args: Iterable[str] = ()) -> None:
+    for task, commands in invocation.tasks.items():
+        prefix = colors.cyan(f"dev-cmd {colors.bold(task)}]")
+        for command in commands:
+            if isinstance(command, Command):
+                await aioconsole.aprint(
+                    f"{prefix} {colors.magenta(f'Executing {colors.bold(command.name)}...')}",
+                    use_stderr=True,
+                )
+                process = await _invoke_command(command, extra_args)
+                returncode = await process.wait()
+                if returncode != 0:
+                    raise CalledProcessError(returncode=returncode, cmd=command.args)
+            else:
+                message = colors.magenta(
+                    f"Parallelizing {len(command)} commands: "
+                    f"{colors.bold(' '.join(cmd.name for cmd in command))}"
+                )
+                await aioconsole.aprint(f"{prefix} {message}...", use_stderr=True)
+                processes = [
+                    await _invoke_command(
+                        cmd,
+                        extra_args,
+                        stdout=asyncio.subprocess.PIPE,
+                        stderr=asyncio.subprocess.STDOUT,
+                    )
+                    for cmd in command
+                ]
+
+                errors: list[tuple[str, CalledProcessError]] = []
+                for cmd, process, (stdout, _) in zip(
+                    command,
+                    processes,
+                    await asyncio.gather(*[process.communicate() for process in processes]),
+                ):
+                    assert process.returncode is not None
+                    if process.returncode != 0:
+                        errors.append(
+                            (
+                                cmd.name,
+                                CalledProcessError(returncode=process.returncode, cmd=cmd.args),
+                            )
+                        )
+                    cmd_name = colors.color(
+                        cmd.name, fg="magenta" if process.returncode == 0 else "red", style="bold"
+                    )
+                    await aioconsole.aprint(f"{prefix} {cmd_name}:", use_stderr=True)
+                    await aioconsole.aprint(stdout.decode(), end="", use_stderr=True)
+                if errors:
+                    lines = [f"{len(errors)} of {len(command)} parallel commands in {task} failed:"]
+                    lines.extend(f"{cmd}: {error}" for cmd, error in errors)
+                    raise ParallelExecutionError(os.linesep.join(lines))
+
+
 def _run(dev: Dev, *tasks: str, extra_args: Iterable[str] = ()) -> None:
     if tasks:
         try:
@@ -57,22 +126,7 @@ def _run(dev: Dev, *tasks: str, extra_args: Iterable[str] = ()) -> None:
             f"arguments: {extra_args}"
         )
 
-    for task, commands in invocation.tasks.items():
-        prefix = colors.cyan(f"dev-cmd {colors.bold(task)}]")
-        for command in commands:
-            print(
-                f"{prefix} {colors.magenta(f'Executing {colors.bold(command.name)}...')}",
-                file=sys.stderr,
-            )
-            args = list(command.args)
-            if extra_args and command.accepts_extra_args:
-                args.extend(extra_args)
-
-            if not os.path.exists(command.cwd):
-                raise InvalidModelError(
-                    f"The `cwd` for command {command.name!r} does not exist: {command.cwd}"
-                )
-            subprocess.run(args, env=command.env, cwd=command.cwd, check=True)
+    return asyncio.run(_invoke(invocation, extra_args))
 
 
 def _parse_args() -> tuple[list[str], list[str]]:
@@ -115,7 +169,7 @@ def main() -> Any:
         return _run(dev, *tasks, extra_args=extra_args)
     except DevError as e:
         return f"{colors.red('Configuration error')}: {colors.yellow(str(e))}"
-    except (OSError, CalledProcessError) as e:
+    except (OSError, CalledProcessError, ParallelExecutionError) as e:
         return colors.red(str(e))
 
 
diff --git a/pyproject.toml b/pyproject.toml
@@ -6,6 +6,7 @@ backend = "setuptools.build_meta"
 name = "dev-cmd"
 requires-python = ">=3.9"
 dependencies = [
+    "aioconsole",
     "ansicolors",
     "tomlkit; python_version < '3.11'",
 ]
@@ -63,6 +64,10 @@ dev = [
     "types-tqdm",
 ]
 
+[[tool.mypy.overrides]]
+module = ["aioconsole.*"]
+follow_untyped_imports = true
+
 [[tool.mypy.overrides]]
 module = ["colors.*"]
 follow_untyped_imports = true
@@ -100,17 +105,30 @@ check-fmt = ["ruff", "format", "--diff"]
 lint = ["ruff", "check", "--fix"]
 check-lint = ["ruff", "check"]
 
-type-check = ["mypy", "dev_cmd"]
-type-check-39 = ["mypy", "--python-version", "3.9", "dev_cmd"]
-type-check-313 = ["mypy", "--python-version",  "3.13", "dev_cmd"]
+type-check = ["mypy", "dev_cmd", "tests"]
+type-check-39 = ["mypy", "--python-version", "3.9", "dev_cmd", "tests"]
+type-check-310 = ["mypy", "--python-version", "3.10", "dev_cmd", "tests"]
+type-check-311 = ["mypy", "--python-version", "3.11", "dev_cmd", "tests"]
+type-check-312 = ["mypy", "--python-version", "3.12", "dev_cmd", "tests"]
+type-check-313 = ["mypy", "--python-version",  "3.13", "dev_cmd", "tests"]
 
 [tool.dev-cmd.commands.test]
 args = ["pytest", "-n", "auto"]
 accepts-extra-args = true
 
 [tool.dev-cmd.aliases]
-checks = ["fmt", "lint", "type-check-39", "type-check-313", "test"]
-ci = ["check-fmt", "check-lint", "type-check", "test"]
+checks = [
+    "fmt",
+    "lint",
+    # Parallelizing the type checks is safe (they don't modify files), and it nets a ~30% speedup
+    # over running them all serially.
+    ["type-check-39", "type-check-310", "type-check-311", "type-check-312", "type-check-313"],
+    "test"
+]
+ci = [
+    ["check-fmt", "check-lint", "type-check"],
+    "test"
+]
 
 [tool.dev-cmd.default]
 alias = "checks"
diff --git a/uv.lock b/uv.lock
