Migrate CLI from manual argv parsing to tyro with component selection

ddss/main.py

@@ -1,24 +1,38 @@
-import sys
 import asyncio
 import tempfile
 import pathlib
+from typing import Annotated, Optional
+import tyro
 from .orm import initialize_database
 from .ds import main as ds
 from .egg import main as egg
 from .input import main as input
 from .output import main as output
+from .load import main as load
+from .dump import main as dump
+
+component_map = {
+    "ds": ds,
+    "egg": egg,
+    "input": input,
+    "output": output,
+    "load": load,
+    "dump": dump,
+}
 
 
-async def main(addr):
+async def run(addr: str, components: list[str]) -> None:
     engine, session = await initialize_database(addr)
+
     try:
+        try:
+            coroutines = [component_map[component](addr, engine, session) for component in components]
+        except KeyError as e:
+            print(f"error: unsupported component {e}")
+            raise asyncio.CancelledError()
+
         await asyncio.wait(
-            [
-                asyncio.create_task(ds(addr, engine, session)),
-                asyncio.create_task(egg(addr, engine, session)),
-                asyncio.create_task(input(addr, engine, session)),
-                asyncio.create_task(output(addr, engine, session)),
-            ],
+            [asyncio.create_task(coro) for coro in coroutines],
             return_when=asyncio.FIRST_COMPLETED,
         )
     except asyncio.CancelledError:
@@ -35,26 +49,46 @@ async def main(addr):
 }
 
 
-def cli():
-    if len(sys.argv) == 1:
+def main(
+    addr: Annotated[
+        Optional[str],
+        tyro.conf.arg(
+            aliases=["-a"],
+            help="Database address URL. If not provided, uses a temporary SQLite database.",
+        ),
+    ] = None,
+    component: Annotated[
+        list[str],
+        tyro.conf.arg(
+            aliases=["-c"],
+            help="Components to run.",
+        ),
+    ] = ["input", "output", "ds", "egg"],
+) -> None:
+    """DDSS - Distributed Deductive System Sorts
+
+    Run DDSS with an interactive deductive environment.
+    """
+    if addr is None:
         tmpdir = tempfile.TemporaryDirectory()
         path = pathlib.Path(tmpdir.name) / "ddss.db"
         addr = f"sqlite:///{path.as_posix()}"
-    elif len(sys.argv) == 2 and sys.argv[1] not in ["--help", "-help", "-h", "/help", "/h", "/?"]:
-        addr = sys.argv[1]
-    else:
-        print(f"Usage: {sys.argv[0]} [<database-addr>]")
-        sys.exit(1)
+
     for key, value in sqlalchemy_driver.items():
         if addr.startswith(f"{key}://"):
             addr = addr.replace(f"{key}://", f"{key}+{value}://")
         if addr.startswith(f"{key}+{value}://"):
             break
     else:
-        print(f"Unsupported database address: {addr}")
-        sys.exit(1)
+        print(f"error: unsupported database: {addr}")
+        raise SystemExit(1)
+
     print(f"addr: {addr}")
-    asyncio.run(main(addr))
+    asyncio.run(run(addr, component))
+
+
+def cli():
+    tyro.cli(main)
 
 
 if __name__ == "__main__":

pyproject.toml

@@ -12,6 +12,7 @@ dependencies = [
     "apyds-egg~=0.0.11",
     "prompt-toolkit~=3.0.52",
     "sqlalchemy[aiosqlite,aiomysql,postgresql-asyncpg]~=2.0.45",
+    "tyro~=1.0.3",
 ]
 
 [project.urls]