pw_cli: Decorator for plugins.Registry registration

Change-Id: I5f6eef65ecc0f6dd6b0b159210ba6533db9a939a
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/40762
Pigweed-Auto-Submit: Wyatt Hepler <hepler@google.com>
Commit-Queue: Auto-Submit <auto-submit@pigweed.google.com.iam.gserviceaccount.com>
Reviewed-by: Keir Mierle <keir@google.com>
Reviewed-by: Rob Mohr <mohrr@google.com>

3 files changed, 93 insertions, 8 deletions

diff --git a/pw_cli/docs.rst b/pw_cli/docs.rst
index 35378e08c..42144f1bd 100644
--- a/pw_cli/docs.rst
+++ b/pw_cli/docs.rst
@@ -289,13 +289,38 @@ registered (see :py:meth:`pw_cli.plugins.Registry.__init__`).
 
 Plugins may be registered in a few different ways.
 
- * Register with a direct function call. See
-   :py:meth:`pw_cli.plugins.Registry.register` and
+ * **Direct function call.** Register plugins by calling
+   :py:meth:`pw_cli.plugins.Registry.register` or
    :py:meth:`pw_cli.plugins.Registry.register_by_name`.
- * Register from plugins files. See
-   :py:meth:`pw_cli.plugins.Registry.register_file` and
-   :py:meth:`pw_cli.plugins.Registry.register_directory`. Plugins files use a
-   simple format:
+
+   .. code-block:: python
+
+     registry = pw_cli.plugins.Registry()
+
+     registry.register('plugin_name', my_plugin)
+     registry.register_by_name('plugin_name', 'module_name', 'function_name')
+
+ * **Decorator.** Register using the :py:meth:`pw_cli.plugins.Registry.plugin`
+   decorator.
+
+   .. code-block:: python
+
+     _REGISTRY = pw_cli.plugins.Registry()
+
+     # This function is registered as the "my_plugin" plugin.
+     @_REGISTRY.plugin
+     def my_plugin():
+         pass
+
+     # This function is registered as the "input" plugin.
+     @_REGISTRY.plugin(name='input')
+     def read_something():
+         pass
+
+   The decorator may be aliased to give a cleaner syntax (e.g. ``register =
+   my_registry.plugin``).
+
+ * **Plugins files.** Plugins files use a simple format:
 
    .. code-block::
 
@@ -304,7 +329,12 @@ Plugins may be registered in a few different ways.
 
      another_plugin some_module some_function
 
-Module reference
-^^^^^^^^^^^^^^^^
+   These files are placed in the file system and apply similarly to Git's
+   ``.gitignore`` files. From Python, these files are registered using
+   :py:meth:`pw_cli.plugins.Registry.register_file` and
+   :py:meth:`pw_cli.plugins.Registry.register_directory`.
+
+pw_cli.plugins module reference
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 .. automodule:: pw_cli.plugins
   :members:
diff --git a/pw_cli/py/plugins_test.py b/pw_cli/py/plugins_test.py
index 908bc4b2e..c2063996d 100644
--- a/pw_cli/py/plugins_test.py
+++ b/pw_cli/py/plugins_test.py
@@ -184,6 +184,27 @@ class TestPluginRegistry(unittest.TestCase):
         finally:
             del sys.modules[fake_module_name]
 
+    def test_decorator_not_called(self) -> None:
+        @self._registry.plugin
+        def nifty() -> None:
+            pass
+
+        self.assertEqual(self._registry['nifty'].target, nifty)
+
+    def test_decorator_called_no_args(self) -> None:
+        @self._registry.plugin()
+        def nifty() -> None:
+            pass
+
+        self.assertEqual(self._registry['nifty'].target, nifty)
+
+    def test_decorator_called_with_args(self) -> None:
+        @self._registry.plugin(name='nifty')
+        def my_nifty_keen_plugin() -> None:
+            pass
+
+        self.assertEqual(self._registry['nifty'].target, my_nifty_keen_plugin)
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/pw_cli/py/pw_cli/plugins.py b/pw_cli/py/pw_cli/plugins.py
index 540b09ca3..d264a39f9 100644
--- a/pw_cli/py/pw_cli/plugins.py
+++ b/pw_cli/py/pw_cli/plugins.py
@@ -34,6 +34,7 @@ import importlib
 import inspect
 import logging
 from pathlib import Path
+import pkgutil
 import sys
 from textwrap import TextWrapper
 import types
@@ -364,6 +365,22 @@ class Registry(collections.abc.Mapping):
         else:
             yield '  (none found)'
 
+    def plugin(self,
+               function: Callable = None,
+               *,
+               name: str = None) -> Callable[[Callable], Callable]:
+        """Decorator that registers a function with this plugin registry."""
+        def decorator(function: Callable) -> Callable:
+            self.register(function.__name__ if name is None else name,
+                          function)
+            return function
+
+        if function is None:
+            return decorator
+
+        self.register(function.__name__, function)
+        return function
+
 
 def find_in_parents(name: str, path: Path) -> Optional[Path]:
     """Searches parent directories of the path for a file or directory."""
@@ -388,3 +405,20 @@ def find_all_in_parents(name: str, path: Path) -> Iterator[Path]:
 
         yield result
         path = result.parent.parent
+
+
+def import_submodules(module: types.ModuleType,
+                      recursive: bool = False) -> None:
+    """Imports the submodules of a package.
+
+    This can be used to collect plugins registered with a decorator from a
+    directory.
+    """
+    path = module.__path__  # type: ignore[attr-defined]
+    if recursive:
+        modules = pkgutil.walk_packages(path, module.__name__ + '.')
+    else:
+        modules = pkgutil.iter_modules(path, module.__name__ + '.')
+
+    for info in modules:
+        importlib.import_module(info.name)