From e055b2064064a0b72aacd73ff416aa0b7702830e Mon Sep 17 00:00:00 2001 From: novirium Date: Sun, 7 Feb 2021 22:04:58 +0800 Subject: [PATCH] Add registration for plugin run and init func --- shepherd/__init__.py | 2 + shepherd/agent/plugin.py | 131 +++++++++++++++++++++++++++---- tests/assets/classtestplugin.py | 12 ++- tests/assets/simpletestplugin.py | 13 +++ tests/test_core.py | 1 + tests/test_plugins.py | 32 +++++--- 6 files changed, 161 insertions(+), 30 deletions(-) diff --git a/shepherd/__init__.py b/shepherd/__init__.py index e6e7c59..edfcd35 100644 --- a/shepherd/__init__.py +++ b/shepherd/__init__.py @@ -3,3 +3,5 @@ from .agent.plugin import plugin_class # noqa from .agent.plugin import plugin_function # noqa from .agent.plugin import plugin_hook # noqa from .agent.plugin import plugin_attachment # noqa +from .agent.plugin import plugin_init # noqa +from .agent.plugin import plugin_run # noqa diff --git a/shepherd/agent/plugin.py b/shepherd/agent/plugin.py index d5c90c0..d0b5adc 100644 --- a/shepherd/agent/plugin.py +++ b/shepherd/agent/plugin.py @@ -212,6 +212,44 @@ def plugin_attachment(hook_identifier): return attachment_decorator +InitMarker = namedtuple("InitMarker", []) + + +def plugin_init(func=None): + """ + Method decorator to register a method as a plugin init function, similar + to passing it to `interface.register_init()`. + + Can either be used on functions in the root level of the plugin module, or + on methods within the registered Plugin Class (either with @plugin_class or + interface.register_plugin_class() ) + """ + if func is None: + return plugin_init + + func._shepherd_load_marker = InitMarker() + return func + + +RunMarker = namedtuple("RunMarker", []) + + +def plugin_run(func=None): + """ + Method decorator to register a method as a plugin run function, similar + to passing it to `interface.register_run()`. + + Can either be used on functions in the root level of the plugin module, or + on methods within the registered Plugin Class (either with @plugin_class or + interface.register_plugin_class() ) + """ + if func is None: + return plugin_run + + func._shepherd_load_marker = RunMarker() + return func + + @preserve.preservable(exclude_attrs=('function')) class InterfaceCall(): def __init__(self, plugin_name, function_name, kwargs=None): @@ -386,6 +424,8 @@ class PluginInterface(): self._tasks = [] self._plugin_class = None self._plugin_obj = None + self._init_func = None + self._run_func = None self.config = None self.plugins = None self.hooks = None @@ -502,6 +542,54 @@ class PluginInterface(): self._hooks[name] = PluginHook(name, signature) return self._hooks[name] + def register_init(self, func): + """ + Register a function or method as the init function for the plugin. This will be called + when the plugin is initialised, after load. This is where the plugin can do any setup + required before hooks and interface functions may be called by other plugins. Plugin config + is available during this call. + + Plugin init is also where any tasks may be added. + + The plugin init function cannot take any arguments. + + The plugin init function is analogous to the `__init__` method when using a plugin class is + registered. If both an init function _and_ a plugin class are registered, both the init + function and the `__init__` method will be called. + """ + self._load_guard() + + if self._init_func is not None: + raise PluginLoadError("Plugin can only register one init function") + if not callable(func): + raise TypeError("Plugin init function must be a callable.") + if len(inspect.signature(func).parameters) > 0: + raise TypeError("Plugin init function cannot take any arguments") + self._init_func = func + + def register_run(self, func): + """ + Register a function or method to be called in a seperate thread once all plugins are + initialised. This function is intended to be used for any continuous loop needed by the + plugin, to avoid blocking other plugins or Shepherd itself. When the "run" function is + called, all other plugin hooks and interface functions are available to be called. + + The plugin "run" function cannot take any arguments. + + If trying to register a method on a plugin class, it is better to use the decorator form + "@plugin_run", as this will then bind to the actual instane of the class once it is + instantiated. + """ + self._load_guard() + + if self._run_func is not None: + raise PluginLoadError("Plugin can only register one run function") + if not callable(func): + raise TypeError("Plugin run function must be a callable.") + if len(inspect.signature(func).parameters) > 0: + raise TypeError("Plugin run function cannot take any arguments") + self._run_func = func + def add_task(self, trigger, interface_function, kwargs=None): """ Add a task when creating a new session. Can only be called during init (object or hook). @@ -713,6 +801,10 @@ def load_plugin_interface(plugin_name, interface, module=None): setattr(module, key, newhook) elif isinstance(attr._shepherd_load_marker, ClassMarker): interface.register_class(attr) + elif isinstance(attr._shepherd_load_marker, InitMarker): + interface.register_init(attr) + elif isinstance(attr._shepherd_load_marker, RunMarker): + interface.register_run(attr) if interface._plugin_class is not None: # Scan plugin class for marked methods @@ -729,6 +821,10 @@ def load_plugin_interface(plugin_name, interface, module=None): # Hooks are a little different in that we replace the attr with the hook newhook = interface.register_hook(**attr._shepherd_load_marker._asdict()) setattr(interface._plugin_class, key, newhook) + elif isinstance(attr._shepherd_load_marker, InitMarker): + interface.register_init(UnboundMethod(attr)) + elif isinstance(attr._shepherd_load_marker, RunMarker): + interface.register_run(UnboundMethod(attr)) # Assemble remote interface function specs @@ -759,21 +855,18 @@ def init_plugins(plugin_configs, plugin_dir=None): plugin_interfaces[plugin_name] = load_plugin(plugin_name, plugin_dir) interface_functions = {} - interface_functions_proxy = MappingProxyType(interface_functions) # Run plugin init and init hooks for plugin_name, interface in plugin_interfaces.items(): - # Though we set `plugins` to the proxy, it's empty until after init - interface.plugins = interface_functions_proxy - interface.config = plugin_configs[plugin_name] + # Collect interface functions from this plugin + interface_functions[plugin_name] = NamespaceProxy(interface._functions) - # TODO This probably should be technically done after all plugins have finished init? Could - # then go in loop that also does the interface functions proxy - interface.hooks = NamespaceProxy(interface._hooks) + # Provide config for plugin init + interface.config = plugin_configs[plugin_name] # If it has one, instantiate the plugin object and bind methods to it. if interface._plugin_class is not None: - # Special case with the 'shepherd' plugin where it is already instantiated + # Special case with the 'shepherd' plugin already populates `_plugin_obj` if interface._plugin_obj is None: interface._plugin_obj = interface._plugin_class() @@ -785,10 +878,15 @@ def init_plugins(plugin_configs, plugin_dir=None): if isinstance(attachment.func, UnboundMethod): attachment.func = attachment.func.bind(interface._plugin_obj) - # TODO Probably need special case for looking for `init` attachments, and run init here, - # before other attachments. Or can we rely on populating `.hooks` later, and that only we - # have access to attachments at this point? Probably. I mean _other_ plugins will have - # attached to this plugin's hooks at this point. + if isinstance(interface._init_func, UnboundMethod): + interface._init_func = interface._init_func.bind(interface._plugin_obj) + + if isinstance(interface._run_func, UnboundMethod): + interface._run_func = interface._run_func.bind(interface._plugin_obj) + + # Call the plugin init func (we've already done any plugin class instance __init__ above) + if interface._init_func is not None: + interface._init_func() # Find hooks attachments are referring to and attach them for attachment in interface._attachments: @@ -806,13 +904,14 @@ def init_plugins(plugin_configs, plugin_dir=None): plugin_interfaces[hook_plugin_name]._hooks[hook_name]._attach(attachment.func, plugin_name) - # TODO We've run the object __init__, but not the init hook, that needs to be done interface._initialised = True # Wait until all plugins have run their init before filling in and giving access - # to all the interface functions. + # to all the interface functions and hooks + interface_functions_proxy = MappingProxyType(interface_functions) for plugin_name, interface in plugin_interfaces.items(): - # Each plugin has a NamespaceProxy of it's interface functions for read-only attr access - interface_functions[plugin_name] = NamespaceProxy(interface._functions) + # Each plugin has a NamespaceProxy of its interface functions for read-only attr access + interface.plugins = interface_functions_proxy + interface.hooks = NamespaceProxy(interface._hooks) return plugin_interfaces diff --git a/tests/assets/classtestplugin.py b/tests/assets/classtestplugin.py index 6f79af6..941183e 100644 --- a/tests/assets/classtestplugin.py +++ b/tests/assets/classtestplugin.py @@ -1,6 +1,7 @@ # pylint: disable=no-self-argument from configspec import * -from shepherd import PluginInterface, plugin_class, plugin_function, plugin_hook, plugin_attachment +from shepherd import PluginInterface, plugin_class, plugin_function, plugin_hook +from shepherd import plugin_attachment, plugin_run, plugin_init """ Plugin to test the plugin class systems and the various decorator markers @@ -36,6 +37,7 @@ class ClassPlugin(): self.interface = interface self.plugins = interface.plugins self.hooks = interface.hooks + self.interface.init_method_called = True # Interface functions @plugin_function @@ -82,3 +84,11 @@ class ClassPlugin(): @staticmethod def static_attach(a, b): return F"static attachment {a} {b}" + + @plugin_init + def plugin_init2_method(self): + self.interface.init2_method_called = True + + @plugin_run + def plugin_run_method(self): + self.interface.run_method_called = True diff --git a/tests/assets/simpletestplugin.py b/tests/assets/simpletestplugin.py index d1667f9..7c75d9f 100644 --- a/tests/assets/simpletestplugin.py +++ b/tests/assets/simpletestplugin.py @@ -40,3 +40,16 @@ interface.register_hook("hook_with_fancy_args", signature(lambda arg_a, arg_b, a interface.register_attachment(basic_attachment, "basic_hook") interface.register_attachment(attachment_with_args, "simpletestplugin.hook_with_args") interface.register_attachment(attachment_with_fancy_args, "hook_with_fancy_args") + + +def my_init_func(): + # Create a dummy variable + interface.init_func_called = True + + +def my_run_func(): + interface.run_func_called = True + + +interface.register_init(my_init_func) +interface.register_run(my_run_func) diff --git a/tests/test_core.py b/tests/test_core.py index 8821a06..07bd0f0 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -89,6 +89,7 @@ def test_local_agent_start(basic_config): def test_local_agent_plugin_start(plugin_config): agent = core.Agent(plugin_config) agent.start() + assert agent.plugin_interfaces["classtestplugin"].run_method_called is True assert agent.interface_functions["classtestplugin"].instance_method( 3) == "instance method 3" diff --git a/tests/test_plugins.py b/tests/test_plugins.py index c2a85df..e4821d5 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -1,5 +1,4 @@ # pylint: disable=redefined-outer-name -import sys from pathlib import Path import pytest @@ -25,7 +24,7 @@ def test_simple_interface_function_load(simple_plugin: plugin.PluginInterface): @pytest.fixture -def simple_running_plugin(request): +def simple_initialised_plugin(request): plugin.unload_plugin("simpletestplugin") interface = plugin.load_plugin("simpletestplugin", Path(request.fspath.dirname)/'assets') # The plugin system is _not_ responsible for making sure the config passed to it is valid. It @@ -39,29 +38,32 @@ def simple_running_plugin(request): return interface -def test_simple_plugin_init(simple_running_plugin): - assert simple_running_plugin._plugin_name == "simpletestplugin" +def test_simple_plugin_init(simple_initialised_plugin): + assert simple_initialised_plugin._plugin_name == "simpletestplugin" + # Check registered init function has run + assert simple_initialised_plugin.init_func_called is True -def test_simple_interface_functions(simple_running_plugin): +def test_simple_interface_functions(simple_initialised_plugin): # Check module level function dict - assert simple_running_plugin._functions["my_interface_function"]() == 42 + assert simple_initialised_plugin._functions["my_interface_function"]() == 42 # Check functions handed back to plugin - assert simple_running_plugin.plugins["simpletestplugin"].my_interface_function() == 42 + assert simple_initialised_plugin.plugins["simpletestplugin"].my_interface_function() == 42 -def test_simple_hook_attachments(simple_running_plugin): - assert "basic_hook" in simple_running_plugin._hooks - assert simple_running_plugin._hooks['basic_hook']() == {'simpletestplugin': "basic attachment"} - assert simple_running_plugin.hooks.hook_with_args( +def test_simple_hook_attachments(simple_initialised_plugin): + assert "basic_hook" in simple_initialised_plugin._hooks + assert simple_initialised_plugin._hooks['basic_hook']( + ) == {'simpletestplugin': "basic attachment"} + assert simple_initialised_plugin.hooks.hook_with_args( 3, 7) == {'simpletestplugin': "attachment with args: 3, 7"} - assert simple_running_plugin.hooks.hook_with_fancy_args( + assert simple_initialised_plugin.hooks.hook_with_fancy_args( 2, 4) == {'simpletestplugin': "attachment with fancy args: 2, 4, True"} with pytest.raises(TypeError, match="takes 2 positional arguments but 3 were"): - simple_running_plugin.hooks.hook_with_args(3, 7, 5) + simple_initialised_plugin.hooks.hook_with_args(3, 7, 5) def test_dirty_plugin_load(request): @@ -85,6 +87,10 @@ def running_class_plugin(request): def test_class_plugin_init(running_class_plugin): assert running_class_plugin._plugin_name == "classtestplugin" + # Check plugin object init method has run + assert running_class_plugin.init_method_called is True + # Check registered init method has run + assert running_class_plugin.init2_method_called is True def test_class_interface_functions(running_class_plugin):