Add registration for plugin run and init func

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

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

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

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)

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"
 

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):