app: main: support custom tools

Add support for loading tools from the custom tools path.

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

Author: JordanYates

src/infuse_iot/app/main.py

@@ -8,13 +8,16 @@
 
 import argparse
 import importlib.util
+import pathlib
 import pkgutil
 import sys
+import types
 
 import argcomplete
 
 import infuse_iot.tools
 from infuse_iot.commands import InfuseCommand
+from infuse_iot.credentials import get_custom_tool_path
 from infuse_iot.version import __version__
 
 
@@ -38,24 +41,42 @@ def run(self, 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,
+        )
+        parser.set_defaults(tool_class=tool_cls)
+        tool_cls.add_parser(parser)
+
     def _load_tools(self, parser: argparse.ArgumentParser):
         tools_parser = parser.add_subparsers(title="commands", metavar="<command>", required=True)
 
-        # Iterate over tools
+        # 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)
-
-            # Add tool to parser
-            tool_cls: InfuseCommand = module.SubCommand
-            parser = tools_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)
+            self._load_from_module(tools_parser, module)
+
+        # 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 _:
+                    continue
+                if hasattr(module, "SubCommand"):
+                    self._load_from_module(tools_parser, module)
 
 
 def main(argv=None):