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>
diff --git a/pw_cli/docs.rst b/pw_cli/docs.rst
index 35378e0..42144f1 100644
--- a/pw_cli/docs.rst
+++ b/pw_cli/docs.rst
@@ -289,13 +289,38 @@
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 @@
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 908bc4b..c206399 100644
--- a/pw_cli/py/plugins_test.py
+++ b/pw_cli/py/plugins_test.py
@@ -184,6 +184,27 @@
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 540b09c..d264a39 100644
--- a/pw_cli/py/pw_cli/plugins.py
+++ b/pw_cli/py/pw_cli/plugins.py
@@ -34,6 +34,7 @@
import inspect
import logging
from pathlib import Path
+import pkgutil
import sys
from textwrap import TextWrapper
import types
@@ -364,6 +365,22 @@
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 @@
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)
