shepherd-agent/shepherd/plugin.py

#!/usr/bin/env python3

from contextlib import suppress
from abc import ABC, abstractmethod
import importlib

from types import SimpleNamespace
from collections import namedtuple
import sys
import os

import shepherd.scheduler


class Hook():
    def __init__(self):
        self.attached_functions = []

    def attach(self, new_func):
        if not callable(new_func):
            raise TypeError("Argument to Hook.attach must be callable")
        self.attached_functions.append(new_func)

    def __call__(self, *args, **kwargs):
        for func in self.attached_functions:
            func(*args, **kwargs)


class InterfaceFunction():
    def __init__(self, func):
        if not callable(func):
            raise TypeError("Argument to InterfaceFunction must be callable")
        self.func = func

    def __call__(self, *args, **kwargs):
        return self.func(*args, **kwargs)


class Plugin(ABC):
    @staticmethod
    @abstractmethod
    def define_config(confdef):
        pass

    @abstractmethod
    def __init__(self, plugininterface, config):
        pass

    def run(self, hooks, plugins):
        pass


class SimplePlugin(Plugin):
    @staticmethod
    def define_config(confdef):
        confdef.add_def()

    def __init__(self, plugininterface, config):
        super().__init__(plugininterface, config)
        self.config = config
        self.interface = plugininterface
        self.plugins = plugininterface.other_plugins
        self.hooks = plugininterface.hooks


plugin_interfaces = {}  # dict of plugin interfaces

# convenience dicts bundling together lists from interfaces
plugin_functions = {}  # dict of plugins containing callable interface functions
plugin_hooks = {}  # dict of plugins containing hook namespaces


_defer = True
_deferred_attachments = []
_deferred_jobs = []


def init_plugins(plugin_classes, plugin_configs, core_config):
    # Startup pluginmanagers
    global plugin_interfaces
    global plugin_functions
    global plugin_hooks

    global _defer
    global _deferred_attachments
    global _deferred_jobs

    for name, plugin_class in plugin_classes.items():
        # Instanciate the plugin interface - this also instanciates
        # the plugin
        plugin_interfaces[name] = PluginInterface(
            name, plugin_class, plugin_configs[name], core_config)
        plugin_functions[name] = plugin_interfaces[name].functions
        plugin_hooks[name] = plugin_interfaces[name].hooks

    # interfaces and hooks should now be populated, attach hooks, schedule jobs
    _defer = False
    for attachment in _deferred_attachments:
        _attach_hook(attachment)

    for job_desc in _deferred_jobs:
        _add_job(job_desc)

    # Hand shared interface callables back out
    for plugininterface in plugin_interfaces.values():
        plugininterface.functions = plugin_functions

def _add_job(job_desc):
    global _deferred_jobs
    global _defer

    if not _defer:
        shepherd.scheduler.schedule_job(job_desc)
    else:
        _deferred_jobs.append(job_desc)

def _attach_hook(attachment):
    global plugin_hooks
    global _deferred_attachments
    global _defer

    if not _defer:
        targetplugin_hooks = plugin_hooks.get(attachment.pluginname)
        if targetplugin_hooks is not None:
            targethook = getattr(targetplugin_hooks, attachment.hookname)
            if targethook is not None:
                targethook.attach(attachment.func)
            else:
                raise Exception("Could not find hook '" +
                                attachment.hookname+"' in module '"+attachment.pluginname+"'")
        else:
            raise Exception(
                "Cannot attach hook to non-existing module '"+attachment.pluginname+"'")
    else:
        _deferred_attachments.append(attachment)

# Eventually, would like to be able to have client plugin simply:
# self.shepherd.add_job(trigger, self.interface.myfunc)
# self.shepherd.attach_hook(pluginnanme,hookname, callable)
# self.shepherd.addinterface
# self.shepherd.hooks.myhook()
# self.shepherd.plugins.otherplugin.otherinterface()

# self.shepherd.add_job()

# Would be good to be able to use abstract methods to enable simpler plugin config
# defs. A way to avoid instantiating the class would be to run it all as class methods,
# enabling


HookAttachment = namedtuple(
    'HookAttachment', ['pluginname', 'hookname', 'func'])


class PluginInterface():
    '''
    Class to handle the management of a single plugin.
    All interaction to or from the plugin to other Shepherd components or
    plugins should go through here.
    '''

    def __init__(self, pluginname, pluginclass, pluginconfig, coreconfig):
        if not issubclass(pluginclass, Plugin):
            raise TypeError(
                "Argument must be subclass of shepherd.plugin.Plugin")

        self.coreconfig = coreconfig

        self.hooks = SimpleNamespace()  # My hooks
        self.functions = SimpleNamespace()  # My callable interface functions

        self._name = pluginname
        self._plugin = pluginclass(self, pluginconfig)

    def register_hook(self, name):
        setattr(self.hooks, name, Hook())

    def register_function(self, func):
        setattr(self.functions, func.__name__, InterfaceFunction(func))

    @property
    def other_plugins(self):
        global plugin_functions
        return plugin_functions

    def attach_hook(self, pluginname, hookname, func):
        _attach_hook(HookAttachment(pluginname, hookname, func))

    # Add a job to the scheduler. By default each will be identified by the interface
    # callable name, and so adding another job with the same callable will oevrride the first.
    # Use the optional job_name to differentiate jobs with an extra string

    def add_job(self, func, hour, minute, second=0, job_name=""):
        for function_name, function in self.functions.__dict__.items():
            if func == function.func:
                # jobstring should canonically describe this job, to be retrieved
                # after reboot later. Of the form:
                # shepherd:pluginname:functionname:jobname
                jobstring = "shepherd:"+self._name+":"+function_name+":"+job_name
                _add_job(shepherd.scheduler.JobDescription(jobstring, hour, minute, second))
                break
        else:
            raise Exception(
                "Could not add job. Callable must first be registered as "
                "a plugin interface with PluginInterface.register_function()")


"""
An interface to a Shepherd module, accessible by other modules.
All public methods in a module interface need to be threadsafe, as they will
be called by other modules (which generally run in a seperate thread)
"""


def find_plugins(plugin_names, plugin_dir=None):
    """
    Looks for the list of plugin names supplied and returns their classes.
    Will first try for plugin modules and packages locally located in ``shepherd.plugins``,
    then for modules and packages prefixed ``shepherd_`` located in the supplied ``plugin_dir`` 
    and lastly in the global import path.
    
    Args:
        plugin_names: List of plugin names to try and load
        plugin_dir: optional search path
    
    Returns:
        Dict of plugin classes, with their names as keys
    """
    plugin_classes = {}
    for plugin_name in plugin_names:
        # First look for core plugins, then the plugin_dir, then in the general import path
        # for custom ones prefixed with "shepherd_"
        try:
            #mod = importlib.import_module("shepherd.plugins." + plugin_name)
            mod = importlib.import_module('.'+plugin_name, "shepherd.plugins")
            #TODO - ModuleNotFoundError is also triggered here if the plugin has a dependancy that can't be found
        except ModuleNotFoundError:
            try:
                if (plugin_dir is not None) and (plugin_dir != ""):
                    if os.path.isdir(plugin_dir):
                        sys.path.append(plugin_dir)
                        mod = importlib.import_module("shepherd_" + plugin_name)
                        sys.path.remove(plugin_dir)
                    else:
                        raise Exception("plugin_dir is not a valid directory")
                else:
                    mod = importlib.import_module("shepherd_" + plugin_name)
            except ModuleNotFoundError:
                raise Exception("Could not find plugin "+plugin_name)

        # Scan imported module for Plugin subclass
        attrs = [getattr(mod, name) for name in dir(mod)]
        for attr in attrs:
            with suppress(TypeError):
                if issubclass(attr, Plugin):
                    plugin_classes[plugin_name] = attr
                    break
        else:
            raise Exception("Imported shepherd plugin modules must contain a "
                            "subclass of shepherd.plugin.Plugin, such as"
                            "shepherd.plugin.SimplePlugin")

    return plugin_classes