app: main: lazy loading of tools

Move tool discovery to a dedicated `registry.py` file, that allows
tools to only be imported when necessary, instead of always importing
all tools.

Assisted-by: GPT-5.5
Signed-off-by: Jordan Yates <jordan@embeint.com>

Author: JordanYates

scripts/custom_tools/custom_tool.py

@@ -14,10 +14,6 @@
 
 
 class SubCommand(InfuseCommand):
-    NAME = "custom_tool"
-    HELP = "Test out-of-tree tool"
-    DESCRIPTION = "Test out-of-tree tool"
-
     @classmethod
     def add_parser(cls, parser):
         parser.add_argument("--echo", "-e", required=True, type=str)

scripts/custom_tools/registry.py

@@ -0,0 +1,14 @@
+#!/usr/bin/env python3
+
+"""Example out-of-tree tool registry."""
+
+from infuse_iot.tools.registry import ToolSpec
+
+TOOLS = (
+    ToolSpec(
+        name="custom_tool",
+        help="Test out-of-tree tool",
+        description="Test out-of-tree tool",
+        module="custom_tool",
+    ),
+)

src/infuse_iot/app/main.py

@@ -7,77 +7,135 @@
 __copyright__ = "Copyright 2024, Embeint Holdings Pty Ltd"
 
 import argparse
+import importlib
 import importlib.util
+import os
 import pathlib
-import pkgutil
 import sys
 import types
+from dataclasses import dataclass
+from typing import Any
 
 import argcomplete
+from argcomplete.lexers import split_line
 
-import infuse_iot.tools
-from infuse_iot.commands import InfuseCommand
 from infuse_iot.credentials import get_custom_tool_path
+from infuse_iot.tools.registry import TOOLS, ToolSpec, load_extension_tools
 from infuse_iot.version import __version__
 
 
+@dataclass(frozen=True)
+class RegisteredTool:
+    spec: ToolSpec
+    extension_path: pathlib.Path | None = None
+
+
 class InfuseApp:
     """The infuse 'application' object"""
 
     def __init__(self):
         self.args = None
         self.parser = argparse.ArgumentParser("infuse")
         self.parser.add_argument("--version", action="version", version=f"{__version__}")
-        self._tools = {}
+        self._tools: dict[str, RegisteredTool] = {}
+        self._tool_parsers = {}
+        self._loaded_tools = set()
         # Load tools
         self._load_tools(self.parser)
-        # Handle CLI tab completion
-        argcomplete.autocomplete(self.parser)
 
     def run(self, argv):
         """Run the chosen subtool handler"""
+        argv = argv or sys.argv[1:]
+        self._load_selected_tool(argv)
+        self._load_selected_completion_tool()
+        # Handle CLI tab completion
+        argcomplete.autocomplete(self.parser)
         self.args = self.parser.parse_args(argv)
 
         tool = self.args.tool_class(self.args)
         tool.run()
 
-    def _load_from_module(self, parent_parser: argparse._SubParsersAction, module: types.ModuleType):
-        tool_cls: InfuseCommand = module.SubCommand
-        parser = parent_parser.add_parser(
-            tool_cls.NAME,
-            help=tool_cls.HELP,
-            description=tool_cls.DESCRIPTION,
-            formatter_class=argparse.RawDescriptionHelpFormatter,
-        )
+    def _load_from_module(
+        self,
+        parent_parser: argparse._SubParsersAction,
+        module: types.ModuleType,
+        parser: argparse.ArgumentParser | None = None,
+    ):
+        tool_cls: Any = module.SubCommand
+        if parser is None:
+            parser = parent_parser.add_parser(
+                tool_cls.NAME,
+                help=tool_cls.HELP,
+                description=tool_cls.DESCRIPTION,
+                formatter_class=argparse.RawDescriptionHelpFormatter,
+            )
         parser.set_defaults(tool_class=tool_cls)
         tool_cls.add_parser(parser)
 
+    def _load_selected_tool(self, argv: list[str]):
+        if not argv:
+            return
+
+        if argv[0] not in self._tools or argv[0] in self._loaded_tools:
+            return
+
+        tool = self._tools[argv[0]]
+        module = self._import_tool_module(tool)
+        self._load_from_module(self._tools_parser, module, self._tool_parsers[tool.spec.name])
+        self._loaded_tools.add(tool.spec.name)
+
+    def _import_tool_module(self, tool: RegisteredTool) -> types.ModuleType:
+        if tool.extension_path is None:
+            return importlib.import_module(tool.spec.module)
+
+        module_path = tool.extension_path / f"{tool.spec.module.replace('.', '/')}.py"
+        module_name = f"infuse_iot_custom_tools.{tool.spec.module}"
+        spec = importlib.util.spec_from_file_location(module_name, module_path)
+        if spec is None or spec.loader is None:
+            raise ImportError(f"Failed to import custom tool module: {module_path}")
+
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        return module
+
+    def _load_selected_completion_tool(self):
+        if "_ARGCOMPLETE" not in os.environ:
+            return
+
+        comp_line = os.environ["COMP_LINE"]
+        comp_point = int(os.environ["COMP_POINT"])
+        _, _, _, comp_words, _ = split_line(comp_line, comp_point)
+
+        # Match argcomplete's own executable/module offset handling.
+        start = int(os.environ["_ARGCOMPLETE"]) - 1
+        parser_words = comp_words[start:]
+        if len(parser_words) > 1:
+            self._load_selected_tool(parser_words[1:])
+
     def _load_tools(self, parser: argparse.ArgumentParser):
-        tools_parser = parser.add_subparsers(title="commands", metavar="<command>", required=True)
+        self._tools_parser = parser.add_subparsers(title="commands", metavar="<command>", required=True)
 
-        # Iterate over local tools
-        for _, name, _ in pkgutil.walk_packages(infuse_iot.tools.__path__):
-            full_name = f"{infuse_iot.tools.__name__}.{name}"
-            module = importlib.import_module(full_name)
-            self._load_from_module(tools_parser, module)
+        # Register local tools without importing their implementation modules.
+        for tool in TOOLS:
+            self._register_tool(tool)
 
         # Load custom tools, if configured
         if extension_tools := get_custom_tool_path():
             extension_path = pathlib.Path(extension_tools)
-            for _, name, _ in pkgutil.walk_packages([extension_tools]):
-                full_name = f"{infuse_iot.tools.__name__}.{name}"
-                full_path = str(extension_path / f"{name}.py")
-                spec = importlib.util.spec_from_file_location(full_name, full_path)
-                if spec is None or spec.loader is None:
-                    continue
-                module = importlib.util.module_from_spec(spec)
-                try:
-                    spec.loader.exec_module(module)
-                except Exception as e:
-                    print(f"Failed to import '{name}': {str(e)}")
-                    continue
-                if hasattr(module, "SubCommand"):
-                    self._load_from_module(tools_parser, module)
+            for tool in load_extension_tools(extension_path):
+                self._register_tool(tool, extension_path)
+
+    def _register_tool(self, tool: ToolSpec, extension_path: pathlib.Path | None = None):
+        if tool.name in self._tools:
+            raise ValueError(f"Tool already registered: {tool.name}")
+
+        self._tools[tool.name] = RegisteredTool(tool, extension_path)
+        self._tool_parsers[tool.name] = self._tools_parser.add_parser(
+            tool.name,
+            help=tool.help,
+            description=tool.description,
+            formatter_class=argparse.RawDescriptionHelpFormatter,
+        )
 
 
 def main(argv=None):

src/infuse_iot/commands.py

@@ -44,21 +44,6 @@ def __init__(self, args: argparse.Namespace):
     def run(self) -> None:
         """Run the subcommand"""
 
-    @property
-    @abstractmethod
-    def NAME(self) -> str:
-        pass
-
-    @property
-    @abstractmethod
-    def HELP(self) -> str:
-        pass
-
-    @property
-    @abstractmethod
-    def DESCRIPTION(self) -> str:
-        pass
-
 
 class InfuseRpcCommand:
     RPC_DATA_SEND: bool = False

src/infuse_iot/tools/annotate_events.py

@@ -45,10 +45,6 @@ class TimeCheckType(enum.Enum):
 
 
 class SubCommand(InfuseCommand):
-    NAME = "annotate_events"
-    HELP = "Annotate events on Infuse Tags"
-    DESCRIPTION = "Save labelled event annotations live on Infuse Tags"
-
     _label_type: LabelType | Path
     _labels: list[str]
     _time_check: TimeCheckType

src/infuse_iot/tools/audio_record.py

@@ -26,10 +26,6 @@
 
 
 class SubCommand(InfuseCommand):
-    NAME = "audio_record"
-    HELP = "Record audio data to a file from TDF"
-    DESCRIPTION = "Record audio data to a file from TDF"
-
     def __init__(self, args):
         self._client = LocalClient(args.server_sock, 1.0)
         self._decoder = TDF()

src/infuse_iot/tools/auto_activate.py

@@ -24,10 +24,6 @@
 
 
 class SubCommand(InfuseCommand):
-    NAME = "auto_activate"
-    HELP = "Automatically activate/deactivate observed devices"
-    DESCRIPTION = "Automatically activate/deactivate observed devices"
-
     def __init__(self, args):
         self.app_ids = args.app
         self.active = args.active or False

src/infuse_iot/tools/bt_log.py

@@ -22,10 +22,6 @@
 
 
 class SubCommand(InfuseCommand):
-    NAME = "bt_log"
-    HELP = "Connect to remote Bluetooth device serial logs"
-    DESCRIPTION = "Connect to remote Bluetooth device serial logs"
-
     def __init__(self, args):
         self._client = LocalClient(args.server_sock, 1.0)
         self._decoder = TDF()

src/infuse_iot/tools/cloud.py

@@ -815,10 +815,6 @@ def upload_diffs_from_application(
 
 
 class SubCommand(InfuseCommand):
-    NAME = "cloud"
-    HELP = "Infuse-IoT cloud interaction"
-    DESCRIPTION = "Infuse-IoT cloud interaction"
-
     @classmethod
     def add_parser(cls, parser):
         parser.add_argument("--api-key", type=str, help="Cloud API key to use instead of stored credentials")

src/infuse_iot/tools/credentials.py

@@ -9,14 +9,11 @@
 
 from infuse_iot import credentials
 from infuse_iot.commands import InfuseCommand
+from infuse_iot.tools.registry import load_extension_tools
 from infuse_iot.util.argparse import ValidDir, ValidFile
 
 
 class SubCommand(InfuseCommand):
-    NAME = "credentials"
-    HELP = "Manage Infuse-IoT credentials"
-    DESCRIPTION = "Manage Infuse-IoT credentials"
-
     @classmethod
     def add_parser(cls, parser):
         parser.add_argument("--api-key", type=str, help="Set Infuse-IoT API key")
@@ -44,6 +41,7 @@ def run(self):
             network_info = yaml.safe_load(content)
             credentials.save_network(network_info["id"], content)
         if self.args.custom_tools:
+            load_extension_tools(self.args.custom_tools)
             credentials.set_custom_tool_path(str(self.args.custom_tools.absolute()))
         if self.args.custom_definitions:
             credentials.set_custom_definitions_path(str(self.args.custom_definitions.absolute()))