From c60d320ebba537e793a40a6bda1d0f3d7a96763c Mon Sep 17 00:00:00 2001
From: Allen Porter <allen.porter@gmail.com>
Date: Tue, 6 Aug 2024 15:07:26 +0000
Subject: [PATCH] Remove boilerplate for new tool setup

---
 home_assistant_datasets/tool/__main__.py      | 41 +++++++++++--------
 .../tool/assist/__init__.py                   | 11 +++--
 .../tool/leaderboard/__init__.py              | 12 ++++--
 3 files changed, 39 insertions(+), 25 deletions(-)

diff --git a/home_assistant_datasets/tool/__main__.py b/home_assistant_datasets/tool/__main__.py
index 81ca283c..7acc96c1 100644
--- a/home_assistant_datasets/tool/__main__.py
+++ b/home_assistant_datasets/tool/__main__.py
@@ -3,36 +3,43 @@
 import argparse
 import importlib
 import logging
+import types
 import sys
 from pathlib import Path
 
-from .leaderboard import prebuild as leaderboard_prebuild, build as leaderboard_build
-from .assist import collect as assist_collect, eval as assist_eval
+from . import assist, leaderboard
 
 
 _LOGGER = logging.getLogger(__name__)
 
+SUBCMDS = {
+    "assist": assist,
+    "leaderboard": leaderboard,
+}
+
+
+def create_arguments(
+    parser: argparse.ArgumentParser, subcmds: dict[str, types.ModuleType]
+) -> None:
+    """Recursively add sub commands/submodules and arguments to the parser"""
+    for name, module in subcmds.items():
+        if not hasattr(module, "SUBCMDS"):
+            module.create_arguments(parser.add_parser(name))
+            continue
+
+        cmd_parser = parser.add_parser(name)
+        subparsers = cmd_parser.add_subparsers(
+            dest="subaction", help="Sub Action", required=True
+        )
+        create_arguments(subparsers, module.SUBCMDS)
+
 
 def get_base_arg_parser() -> argparse.ArgumentParser:
     """Get a base argument parser."""
     parser = argparse.ArgumentParser(description="Home Assistant Datasets Utility")
     parser.add_argument("--debug", action="store_true", help="Enable log output")
     subparsers = parser.add_subparsers(dest="action", help="Action", required=True)
-
-    # Subcommands
-    assist_parser = subparsers.add_parser("assist")
-    assist_subparsers = assist_parser.add_subparsers(
-        dest="subaction", help="Sub Action", required=True
-    )
-    assist_collect.create_arguments(assist_subparsers.add_parser("collect"))
-    assist_eval.create_arguments(assist_subparsers.add_parser("eval"))
-
-    leaderboard_parser = subparsers.add_parser("leaderboard")
-    leaderboard_subparsers = leaderboard_parser.add_subparsers(
-        dest="subaction", help="Sub Action", required=True
-    )
-    leaderboard_prebuild.create_arguments(leaderboard_subparsers.add_parser("prebuild"))
-    leaderboard_build.create_arguments(leaderboard_subparsers.add_parser("build"))
+    create_arguments(subparsers, SUBCMDS)
 
     return parser
 
diff --git a/home_assistant_datasets/tool/assist/__init__.py b/home_assistant_datasets/tool/assist/__init__.py
index 9976d4df..aceb82ae 100644
--- a/home_assistant_datasets/tool/assist/__init__.py
+++ b/home_assistant_datasets/tool/assist/__init__.py
@@ -11,7 +11,10 @@
 ```
 """
 
-__all__ = [
-    "collect",
-    "eval",
-]
+from . import collect, eval
+
+SUBCMDS = {
+    "collect": collect,
+    "eval": eval,
+}
+__all__ = list(SUBCMDS)
diff --git a/home_assistant_datasets/tool/leaderboard/__init__.py b/home_assistant_datasets/tool/leaderboard/__init__.py
index 08e37a62..619b1c73 100644
--- a/home_assistant_datasets/tool/leaderboard/__init__.py
+++ b/home_assistant_datasets/tool/leaderboard/__init__.py
@@ -11,7 +11,11 @@
 ```
 """
 
-__all__ = [
-    "prebuild",
-    "build",
-]
+from . import prebuild, build
+
+
+SUBCMDS = {
+    "prebuild": prebuild,
+    "build": build,
+}
+__all__ = list(SUBCMDS)