Merge pull request #237 from ezmsg-org/chore/cmdline-tlc

QOL: Command Line TLC

Author: griffinmilsap

pyproject.toml

@@ -55,7 +55,6 @@ axisarray = [
 
 [project.scripts]
 ezmsg = "ezmsg.core.command:cmdline"
-ezmsg-perf = "ezmsg.util.perf.command:command"
 
 [project.optional-dependencies]
 axisarray = [

src/ezmsg/core/command.py

@@ -1,33 +1,30 @@
 import argparse
 import asyncio
-import base64
-import json
-import logging
-import subprocess
-import sys
-import webbrowser
-import zlib
+import inspect
 
-from .graphserver import GraphService
+from ezmsg.util.perf.command import setup_perf_cmdline
+
+from .commands import setup_core_cmdline
+from .commands.graphviz import handle_graphviz
+from .commands.mermaid import handle_mermaid, mermaid_url as mm
+from .commands.serve import handle_serve
+from .commands.shutdown import handle_shutdown
+from .commands.start import handle_start
 from .netprotocol import (
     Address,
     GRAPHSERVER_ADDR_ENV,
     GRAPHSERVER_PORT_DEFAULT,
     PUBLISHER_START_PORT_ENV,
     PUBLISHER_START_PORT_DEFAULT,
-    close_stream_writer,
 )
 
-logger = logging.getLogger("ezmsg")
-
 
-def cmdline() -> None:
+def build_parser() -> argparse.ArgumentParser:
     """
-    Command-line interface for ezmsg core server management.
+    Build the ezmsg core command-line parser.
 
-    Provides commands for starting, stopping, and managing ezmsg server
-    processes including GraphServer and SHMServer, as well as utilities
-    for graph visualization.
+    Each command gets its own subparser so command-specific options are not
+    shared globally across unrelated commands.
     """
     parser = argparse.ArgumentParser(
         "ezmsg.core",
@@ -38,63 +35,27 @@ def cmdline() -> None:
             Publishers will be assigned available ports starting from {PUBLISHER_START_PORT_DEFAULT}. (Change with ${PUBLISHER_START_PORT_ENV})
         """,
     )
+    subparsers = parser.add_subparsers(dest="command", required=True, help="command for ezmsg")
 
-    parser.add_argument(
-        "command",
-        help="command for ezmsg",
-        choices=["serve", "start", "shutdown", "graphviz", "mermaid"],
-    )
-
-    parser.add_argument("--address", help="Address for GraphServer", default=None)
-
-    parser.add_argument(
-        "--target",
-        help="Target for mermaid output. Options are 'ink', 'live', and 'play'.",
-        default="live",
-    )
-
-    parser.add_argument(
-        "-c",
-        "--compact",
-        help="""Use compact graph representation. Only used when `cmd` is 'mermaid' or 'graphviz'.
-        Removes the lowest level of detail (typically streams). Can be stacked (eg. '-cc').
-        Warning: this will also prune the graph of proxy topics (nodes that are both sources and targets).
-        """,
-        action="count",
-    )
-
-    parser.add_argument(
-        "-n",
-        "--nobrowser",
-        help="Do not automatically open the browser for mermaid output. `--target` value will be ignored.",
-        action="store_true",
-    )
-
-    class Args:
-        command: str
-        address: str | None
-        target: str
-        compact: int | None
-        nobrowser: bool
+    setup_core_cmdline(subparsers)
+    setup_perf_cmdline(subparsers)
+    return parser
 
-    args = parser.parse_args(namespace=Args)
 
-    graph_address = Address("127.0.0.1", GRAPHSERVER_PORT_DEFAULT)
-    if args.address is not None:
-        graph_address = Address.from_string(args.address)
+def cmdline(argv: list[str] | None = None) -> None:
+    """
+    Command-line interface for ezmsg core server management.
 
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
+    Provides commands for starting, stopping, and managing ezmsg server
+    processes including GraphServer and SHMServer, as well as utilities
+    for graph visualization.
+    """
+    parser = build_parser()
+    args = parser.parse_args(args=argv)
 
-    loop.run_until_complete(
-        run_command(
-            args.command,
-            graph_address,
-            args.target,
-            args.compact,
-            args.nobrowser,
-        )
-    )
+    result = args._handler(args)
+    if inspect.isawaitable(result):
+        asyncio.run(result)
 
 
 async def run_command(
@@ -104,122 +65,20 @@ async def run_command(
     compact: int | None = None,
     nobrowser: bool = False,
 ) -> None:
-    """
-    Run an ezmsg command with the specified parameters.
-
-    This function handles various ezmsg commands like 'serve', 'start', 'shutdown', etc.
-    and manages the graph and shared memory services.
-
-    :param cmd: The command to execute ('serve', 'start', 'shutdown', 'graphviz', 'mermaid')
-    :type cmd: str
-    :param graph_address: Address of the graph service
-    :type graph_address: Address
-    :param target: Target for visualization commands (default: 'live')
-    :type target: str
-    :param compact: Compactification level for visualization commands
-    :type compact: int | None
-    :param nobrowser: Whether to suppress browser opening for visualization
-    :type nobrowser: bool
-    """
-    graph_service = GraphService(graph_address)
-
-    if cmd == "serve":
-        logger.info(f"GraphServer Address: {graph_address}")
-
-        graph_server = graph_service.create_server()
-
-        try:
-            logger.info("Servers running...")
-            graph_server.join()
-
-        except KeyboardInterrupt:
-            logger.info("Interrupt detected; shutting down servers")
-
-        finally:
-            if graph_server is not None:
-                graph_server.stop()
-
-    elif cmd == "start":
-        popen = subprocess.Popen(
-            [sys.executable, "-m", "ezmsg.core", "serve", f"--address={graph_address}"]
-        )
-
-        while True:
-            try:
-                _, writer = await graph_service.open_connection()
-                await close_stream_writer(writer)
-                break
-            except ConnectionRefusedError:
-                await asyncio.sleep(0.1)
-
-        logger.info(f"Forked ezmsg servers in PID: {popen.pid}")
-
-    elif cmd == "shutdown":
-        try:
-            await graph_service.shutdown()
-            logger.info(
-                f"Issued shutdown command to GraphServer @ {graph_service.address}"
-            )
-
-        except ConnectionRefusedError:
-            logger.warning(
-                f"Could not issue shutdown command to GraphServer @ {graph_service.address}; server not running?"
-            )
-
-    elif cmd in ["graphviz", "mermaid"]:
-        graph_out = await graph_service.get_formatted_graph(
-            fmt=cmd, compact_level=compact
-        )
-        print(graph_out)
-        if cmd == "mermaid":
-            if not nobrowser:
-                if target == "live":
-                    print(
-                        "%% If the graph does not render immediately, try toggling the 'Pan & Zoom' button."
-                    )
-                webbrowser.open(mm(graph_out, target=target))
-
-
-def mm(graph: str, target="live") -> str:
-    """
-    Generate a Mermaid visualization URL for the given graph.
-
-    :param graph: Graph representation string to visualize.
-    :type graph: str
-    :param target: Target platform ('live' or 'ink').
-    :type target: str
-    :return: URL for graph visualization.
-    :rtype: str
-    """
-    if target != "ink":
-        jdict = {
-            "code": graph,
-            "mermaid": {"theme": "default"},
-            "updateDiagram": True,
-            "autoSync": True,
-            "rough": False,
-        }
-        graph = json.dumps(jdict)
-    graphbytes: bytes = graph.encode("utf8")
-
-    if target != "ink":
-        compress = zlib.compressobj(9, zlib.DEFLATED, 15, 8, zlib.Z_DEFAULT_STRATEGY)
-        graphbytes = compress.compress(graphbytes)
-        graphbytes += compress.flush()
-
-    base64_bytes = base64.b64encode(graphbytes)
-    base64_string = base64_bytes.decode("ascii")
-
-    if target == "ink":
-        prefix = "https://mermaid.ink/img/"
-    elif target in ["live", "play"]:
-        type_str = "pako"  # or "base64" if we skip compression above.
-        if target == "live":
-            prefix = f"https://mermaid.live/edit#{type_str}:"
-        else:  # "play"
-            prefix = f"https://www.mermaidchart.com/play#{type_str}:"
-    else:
-        raise ValueError(
-            f"Unknown mermaid target '{target}'. Available options are 'ink', 'live', or 'play'."
-        )
-    return prefix + base64_string
+    handlers = {
+        "serve": handle_serve,
+        "start": handle_start,
+        "shutdown": handle_shutdown,
+        "graphviz": handle_graphviz,
+        "mermaid": handle_mermaid,
+    }
+    if cmd not in handlers:
+        raise ValueError(f"Unknown ezmsg command '{cmd}'")
+    args = argparse.Namespace(
+        command=cmd,
+        address=str(graph_address),
+        target=target,
+        compact=compact,
+        nobrowser=nobrowser,
+    )
+    await handlers[cmd](args)

src/ezmsg/core/commands/__init__.py

@@ -0,0 +1,15 @@
+import argparse
+
+from .graphviz import setup_graphviz_cmdline
+from .mermaid import setup_mermaid_cmdline
+from .serve import setup_serve_cmdline
+from .shutdown import setup_shutdown_cmdline
+from .start import setup_start_cmdline
+
+
+def setup_core_cmdline(subparsers: argparse._SubParsersAction) -> None:
+    setup_serve_cmdline(subparsers)
+    setup_start_cmdline(subparsers)
+    setup_shutdown_cmdline(subparsers)
+    setup_graphviz_cmdline(subparsers)
+    setup_mermaid_cmdline(subparsers)

src/ezmsg/core/commands/common.py

@@ -0,0 +1,25 @@
+import argparse
+
+from ..netprotocol import Address, GRAPHSERVER_PORT_DEFAULT
+
+
+def add_address_argument(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument("--address", help="Address for GraphServer", default=None)
+
+
+def add_compact_argument(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument(
+        "-c",
+        "--compact",
+        help="""Use compact graph representation.
+        Removes the lowest level of detail (typically streams). Can be stacked (eg. '-cc').
+        Warning: this will also prune the graph of proxy topics (nodes that are both sources and targets).
+        """,
+        action="count",
+    )
+
+
+def graph_address_from_args(args: argparse.Namespace) -> Address:
+    if args.address is None:
+        return Address("127.0.0.1", GRAPHSERVER_PORT_DEFAULT)
+    return Address.from_string(args.address)

src/ezmsg/core/commands/graphviz.py

@@ -0,0 +1,19 @@
+import argparse
+
+from ..graphserver import GraphService
+from .common import add_address_argument, add_compact_argument, graph_address_from_args
+
+
+async def handle_graphviz(args: argparse.Namespace) -> None:
+    graph_service = GraphService(graph_address_from_args(args))
+    graph_out = await graph_service.get_formatted_graph(
+        fmt="graphviz", compact_level=args.compact
+    )
+    print(graph_out)
+
+
+def setup_graphviz_cmdline(subparsers: argparse._SubParsersAction) -> None:
+    parser = subparsers.add_parser("graphviz")
+    add_address_argument(parser)
+    add_compact_argument(parser)
+    parser.set_defaults(_handler=handle_graphviz)