From b08a93476fff1f9d47aca1952e03ec23a22a10ed Mon Sep 17 00:00:00 2001 From: Christopher Sams Date: Fri, 30 Nov 2018 12:56:51 -0600 Subject: [PATCH] Refactor plugin loader to allow eggs and custom plugin locations. The idea is to allow users to write their own plugins without needing to download the code. They can pip install or dnf install and use the config file to point to their custom plugins. * Refactor the loading logic to use importlib and pkgutil so it can load plugins from an egg and arbitrary places on the filesystem. * Add a custom_plugins option to the general section of the config. Users can add a space separated list of directories where custom plugins reside. Example: ``` [general] email = Joe User custom_plugins = ~/.did/custom_plugins [thing] type = thing ``` --- did/base.py | 7 +++ did/cli.py | 13 +++- did/plugins/__init__.py | 110 +------------------------------- did/stats.py | 63 ++++++++++++++++--- did/utils.py | 117 ++++++++++++++++++++++++++++++++++- tests/plugins/test_detect.py | 15 ----- tests/test_utils.py | 32 ++++++++++ 7 files changed, 223 insertions(+), 134 deletions(-) delete mode 100644 tests/plugins/test_detect.py diff --git a/did/base.py b/did/base.py index 9a8ce5d3..5b80e335 100644 --- a/did/base.py +++ b/did/base.py @@ -108,6 +108,13 @@ def __init__(self, config=None, path=None): raise ConfigFileError( "Unable to read the config file '{0}'.".format(path)) + @property + def custom_plugins(self): + try: + return self.parser.get("general", "custom_plugins") + except: + return None + @property def email(self): """ User email(s) """ diff --git a/did/cli.py b/did/cli.py index 6c5b6b86..77135476 100644 --- a/did/cli.py +++ b/did/cli.py @@ -15,9 +15,9 @@ from dateutil.relativedelta import relativedelta as delta import did.base -import did.utils as utils -from did.utils import log +from did import utils from did.stats import UserStats +from did.utils import log # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Options @@ -173,6 +173,13 @@ def main(arguments=None): """ try: # Parse options, initialize gathered stats + utils.load_components("did.plugins", continue_on_error=True) + + custom_plugins = did.base.Config().custom_plugins + if custom_plugins: + custom_plugins = [p.strip() for p in utils.split(custom_plugins)] + utils.load_components(*custom_plugins, continue_on_error=True) + options, header = Options(arguments).parse() gathered_stats = [] @@ -208,7 +215,7 @@ def main(arguments=None): # Return all gathered stats objects return gathered_stats, team_stats - except did.base.ConfigFileError as error: + except did.base.ConfigFileError: utils.info("Create at least a minimum config file {0}:\n{1}".format( did.base.Config.path(), did.base.Config.example().strip())) raise diff --git a/did/plugins/__init__.py b/did/plugins/__init__.py index 2662148d..fd932a52 100644 --- a/did/plugins/__init__.py +++ b/did/plugins/__init__.py @@ -1,9 +1,9 @@ """ -Modules in this direcotry are searched for available stats. Each +Modules in this directory are searched for available stats. Each plugin should contain a single class inheriting from StatsGroup. Stats from this group will be included in the report if enabled in user config. Name of the plugin should match config section type. -Attribute ``order`` defines the default order in the final report. +Attribute ``order`` defines the order in the final report. This is the default plugin order: @@ -46,110 +46,4 @@ +----------+-----+ | footer | 900 | +----------+-----+ - """ - -from __future__ import unicode_literals, absolute_import - -import os -import sys -import types - -from did.utils import log -from did.base import Config, ConfigError -from did.stats import StatsGroup, EmptyStatsGroup - -# Self reference and file path to this module -PLUGINS = sys.modules[__name__] -PLUGINS_PATH = os.path.dirname(PLUGINS.__file__) - -FAILED_PLUGINS = [] - - -def load(): - """ Check available plugins and attempt to import them """ - # Code is based on beaker-client's command.py script - plugins = [] - for filename in os.listdir(PLUGINS_PATH): - if not filename.endswith(".py") or filename.startswith("_"): - continue - if not os.path.isfile(os.path.join(PLUGINS_PATH, filename)): - continue - plugin = filename[:-3] - if plugin in FAILED_PLUGINS: - # Skip loading plugins that already failed before - continue - try: - __import__(PLUGINS.__name__, {}, {}, [plugin]) - plugins.append(plugin) - log.debug("Successfully imported {0} plugin".format(plugin)) - except (ImportError, SyntaxError) as error: - # Give a warning only when the plugin is configured - message = "Failed to import {0} plugin ({1})".format(plugin, error) - if Config().sections(kind=plugin): - log.warn(message) - else: - log.debug(message) - FAILED_PLUGINS.append(plugin) - return plugins - - -def detect(): - """ - Detect available plugins and return enabled/configured stats - - Yields tuples of the form (section, statsgroup) sorted by the - default StatsGroup order which maybe overriden in the config - file. The 'section' is the name of the configuration section - as well as the option used to enable those particular stats. - """ - - # Load plugins and config - plugins = load() - config = Config() - - # Make sure that all sections have a valid plugin type defined - for section in config.sections(): - if section == 'general': - continue - try: - type_ = config.item(section, 'type') - except ConfigError: - raise ConfigError( - "Plugin type not defined in section '{0}'.".format(section)) - if type_ not in plugins: - raise ConfigError( - "Invalid plugin type '{0}' in section '{1}'.".format( - type_, section)) - - # Detect classes inherited from StatsGroup and return them sorted - stats = [] - for plugin in plugins: - module = getattr(PLUGINS, plugin) - for object_name in dir(module): - statsgroup = getattr(module, object_name) - # Filter out anything except for StatsGroup descendants - if (not isinstance(statsgroup, (type, types.ClassType)) - or not issubclass(statsgroup, StatsGroup) - or statsgroup is StatsGroup - or statsgroup is EmptyStatsGroup): - continue - # Search config for sections with type matching the plugin, - # use order provided there or class default otherwise - for section in config.sections(kind=plugin): - try: - order = int(config.item(section, "order")) - except ConfigError: - order = statsgroup.order - except ValueError: - log.warn("Invalid {0} stats order: '{1}'".format( - section, config.item(section, "order"))) - order = statsgroup.order - stats.append((section, statsgroup, order)) - log.info("Found {0}, an instance of {1}, order {2}".format( - section, statsgroup.__name__, order)) - # Custom stats are handled with a single instance - if statsgroup.__name__ == "CustomStats": - break - for section, statsgroup, _ in sorted(stats, key=lambda x: x[2]): - yield section, statsgroup diff --git a/did/stats.py b/did/stats.py index f56b729a..a693df51 100644 --- a/did/stats.py +++ b/did/stats.py @@ -114,9 +114,37 @@ def merge(self, other): # Stats Group # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +class StatsGroupPlugin(type): + registry = {} + ignore = set([ + "StatsGroupPlugin", + "StatsGroup", + "EmptyStatsGroup", + "UserStats", + ]) + + def __init__(cls, name, bases, attrs): + if name in StatsGroupPlugin.ignore: + return + + plugin_name = cls.__module__.split(".")[-1] + registry = StatsGroupPlugin.registry + + if plugin_name in registry: + orig = registry[plugin_name] + log.warn("%s overriding %s" % (cls.__module__, orig.__module__)) + + registry[plugin_name] = cls + + +# For python3, remove __metaclass__ and redefine as: +# class StatsGroup(six.with_metaclass(StatsGroupPlugin, Stats)): class StatsGroup(Stats): """ Stats group """ + # autoregister all subclasses + __metaclass__ = StatsGroupPlugin + # Default order order = 500 @@ -157,17 +185,13 @@ def fetch(self): class UserStats(StatsGroup): """ User statistics in one place """ - def __init__(self, user=None, options=None): + def __init__(self, user=None, options=None, config=None): """ Initialize stats objects. """ super(UserStats, self).__init__( option="all", user=user, options=options) - self.stats = [] + config = config or did.base.Config() try: - import did.plugins - for section, statsgroup in did.plugins.detect(): - self.stats.append(statsgroup( - option=section, parent=self, - user=self.user.clone(section) if self.user else None)) + self.stats = self.configured_plugins(config) except did.base.ConfigFileError as error: # Missing config file is OK if building options (--help). # Otherwise raise the expection to suggest config example. @@ -177,6 +201,31 @@ def __init__(self, user=None, options=None): else: raise + def configured_plugins(self, cfg): + results = [] + for section in cfg.sections(): + if section == "general": + continue + + data = dict(cfg.section(section, skip=set())) + type_ = data.get("type") + + if not type_: + msg = "Plugin type not defined in section '{0}'." + raise did.base.ConfigError(msg.format(section)) + + if type_ not in StatsGroupPlugin.registry: + raise did.base.ConfigError( + "Invalid plugin type '{0}' in section '{1}'.".format( + type_, section)) + + user = self.user.clone(section) if self.user else None + statsgroup = StatsGroupPlugin.registry[type_] + obj = statsgroup(option=section, parent=self, user=user) + obj.order = data.get("order", statsgroup.order) + results.append(obj) + return sorted(results, key=lambda x: x.order) + def add_option(self, parser): """ Add options for each stats group. """ for stat in self.stats: diff --git a/did/utils.py b/did/utils.py index e9d96c5a..d66a38b5 100644 --- a/did/utils.py +++ b/did/utils.py @@ -4,10 +4,12 @@ from __future__ import unicode_literals, absolute_import +import importlib +import logging import os +import pkgutil import re import sys -import logging import unicodedata from pprint import pformat as pretty @@ -39,6 +41,119 @@ # Utils # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +def _find_base(path): + """ + Given a path to a python file or package, find the top level directory + that isn't a valid package. + """ + if not os.path.isdir(path): + path = os.path.dirname(path) + + if not os.path.exists(os.path.join(path, "__init__.py")): + return None + + while os.path.exists(os.path.join(path, "__init__.py")): + path = os.path.dirname(path) + + return path + + +def _import(path, continue_on_error): + """ Eats or raises import exceptions based on ``continue_on_error``. """ + log.debug("Importing %s" % path) + try: + # importlib is available in stdlib from 2.7+ + return importlib.import_module(path) + except Exception as ex: + log.exception(ex) + if not continue_on_error: + raise + + +def _load_components(path, include=".*", exclude="test", + continue_on_error=True): + num_loaded = 0 + if path.endswith(".py"): + path, _ = os.path.splitext(path) + + path = path.rstrip("/").replace("/", ".") + + package = _import(path, continue_on_error) + if not package: + return 0 + + num_loaded += 1 + + do_include = re.compile(include).search if include else lambda x: True + do_exclude = re.compile(exclude).search if exclude else lambda x: False + + if not hasattr(package, "__path__"): + return num_loaded + + prefix = package.__name__ + "." + for _, name, is_pkg in pkgutil.iter_modules(path=package.__path__, + prefix=prefix): + if not name.startswith(prefix): + name = prefix + name + if is_pkg: + num_loaded += _load_components(name, include, exclude, + continue_on_error) + else: + if do_include(name) and not do_exclude(name): + _import(name, continue_on_error) + num_loaded += 1 + + return num_loaded + + +def load_components(*paths, **kwargs): + """ + Loads all components on the paths. Each path should be a package or module. + All components beneath a path are loaded. This method works whether the + package or module is on the filesystem or in an .egg. If it's in an egg, + the egg must already be on the ``PYTHONPATH``. + + Args: + paths (str): A package or module to load + + Keyword Args: + include (str): A regular expression of packages and modules to include. + Defaults to '.*' + exclude (str): A regular expression of packges and modules to exclude. + Defaults to 'test' + continue_on_error (bool): If True, continue importing even if something + raises an ImportError. If False, raise the first ImportError. + + Returns: + int: The total number of modules loaded. + + Raises: + ImportError + """ + continue_on_error = kwargs.get("continue_on_error", True) + num_loaded = 0 + for path in paths: + tmp = os.path.expandvars(os.path.expanduser(path)) + fs_path = os.path.realpath(tmp) + if os.path.exists(fs_path): + base = _find_base(fs_path) + if not base: + msg = "%s is not a valid python module or package." % path + if continue_on_error: + log.info(msg) + continue + else: + raise ImportError(path) + if base not in sys.path: + sys.path.insert(0, base) + + target = os.path.relpath(fs_path, base) + num_loaded += _load_components(target, **kwargs) + else: + num_loaded += _load_components(path, **kwargs) + return num_loaded + + def eprint(text): """ Print (optionaly encoded) text """ # When there's no terminal we need to explicitly encode strings. diff --git a/tests/plugins/test_detect.py b/tests/plugins/test_detect.py deleted file mode 100644 index ff5dc464..00000000 --- a/tests/plugins/test_detect.py +++ /dev/null @@ -1,15 +0,0 @@ -#!/usr/bin/env python -# -*- coding: utf-8 -*- -# Author: "Chris Ward" - -from __future__ import unicode_literals, absolute_import - - -def test_load(): - from did.plugins import load - assert load - - -def test_detect(): - from did.plugins import detect - assert detect diff --git a/tests/test_utils.py b/tests/test_utils.py index 7b8b7c37..991592b0 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -3,6 +3,10 @@ # Author: "Chris Ward" from __future__ import unicode_literals, absolute_import +import pytest +import os + +import did def test_utils_import(): @@ -54,6 +58,34 @@ def test_log(): # Utils # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +def test_import_success(): + import sys + from did.utils import _import + s = _import("sys", True) + assert s is sys + + s = _import("blah", True) + assert s is None + + +def test_find_base(): + top = os.path.dirname(os.path.dirname(did.__file__)) + assert top == did.utils._find_base(__file__) + assert top == did.utils._find_base(os.path.dirname(__file__)) + + +def test_load_components(): + top = os.path.dirname(did.__file__) + assert did.utils.load_components(top) > 0 + assert did.utils.load_components("did.plugins") > 0 + + +def test_import_failure(): + from did.utils import _import + with pytest.raises(ImportError): + _import("blah", False) + + def test_header(): from did.utils import header assert header