From d6cc2983ea58b398021a5d75025ae8d100f0b776 Mon Sep 17 00:00:00 2001 From: novirium Date: Sat, 4 Jan 2020 16:22:57 +0800 Subject: [PATCH] Add config template generation --- shepherd/core.py | 82 ++++++++++++++++++++++++++++++++++++++++++++-- shepherd/plugin.py | 16 +++++---- 2 files changed, 90 insertions(+), 8 deletions(-) diff --git a/shepherd/core.py b/shepherd/core.py index ecf4585..2cf2ab4 100644 --- a/shepherd/core.py +++ b/shepherd/core.py @@ -1,5 +1,5 @@ """ -Core shepherd module, tying together main service functionality. +Core shepherd module, tying together main service functionality. Provides main CLI. """ @@ -41,7 +41,10 @@ def cli(ctx, default_config_path, local_operation, only_default_layer): ending with ".toml" will be used. """ version_text = pkg_resources.get_distribution("shepherd") - log.info(F"Initialising Shepherd Agent {version_text}") + log.info(F"Shepherd Agent {version_text}") + + if ctx.invoked_subcommand == "template": + return if not default_config_path: default_config_path = sorted(glob.glob("./shepherd*.toml"))[:1] @@ -91,6 +94,81 @@ def test(): print("test!") +@cli.command() +@click.argument('plugin_name', required=False) +@click.option('-a', '--include-all', is_flag=True, + help="Include all optional fields in the template") +@click.option('-c', '--config', 'config_path', type=click.Path(), + help="Path to append or create config tempalate") +@click.option('-d', '--plugin-dir', type=click.Path(), + help="Directory to search for plugin modules, in addition to built in Shepherd" + " plugins and the global import path. Defaults to current directory.") +@click.pass_context +def template(ctx, plugin_name, include_all, config_path, plugin_dir): + """ + Generate a template config TOML file for PLUGIN, or for the Shepherd core if + PLUGIN is not provided. + + If config path is provided ("-c"), append to that file (if it exists) or write to + a new file (if it doesn't yet exist). + """ + + if not plugin_dir: + plugin_dir = Path.cwd() + + confspec = ConfigSpecification() + if (not plugin_name) or (plugin_name=="shepherd"): + plugin_name = "shepherd" + define_core_config(confspec) + else: + try: + plugin_class = plugin.find_plugins([plugin_name], plugin_dir)[0] + except plugin.PluginLoadError as e: + log.error(e.args[0]) + sys.exit(1) + plugin_class.define_config(confspec) + + template_dict = confspec.get_template(include_all) + template_toml = toml.dumps({plugin_name: template_dict}) + + log.info(F"Config template for [{plugin_name}]: \n\n"+template_toml) + + if not config_path: + # reuse parent "-c" for convenience + config_path = ctx.parent.params["default_config_path"] + + if not config_path: + return + + if Path(config_path).is_file(): + try: + existing_config = toml.load(config_path) + except Exception: + click.confirm( + F"File {config_path} already exists and is not a valid TOML file. Overwrite?", + default=True, abort=True) + + click.echo(F"Writing [{plugin_name}] template to {config_path}") + with open(config_path, 'w+') as f: + f.write(template_toml) + else: + if plugin_name in existing_config: + click.confirm(F"Overwrite [{plugin_name}] section in {config_path}?", + default=True, abort=True) + click.echo(F"Overwriting [{plugin_name}] section in {config_path}") + else: + click.confirm(F"Add [{plugin_name}] section to {config_path}?", + default=True, abort=True) + click.echo(F"Adding [{plugin_name}] section to {config_path}") + existing_config[plugin_name] = template_dict + with open(config_path, 'w+') as f: + f.write(toml.dumps(existing_config)) + else: + click.echo(F"Writing [{plugin_name}] template to {config_path}") + with open(config_path, 'w+') as f: + f.write(template_toml) + + def compile_config_and_get_plugins(confman, default_config_path, layers_disabled): """ Run through the process of assembling the various config layers, falling back to working diff --git a/shepherd/plugin.py b/shepherd/plugin.py index ecdd125..d02df0f 100644 --- a/shepherd/plugin.py +++ b/shepherd/plugin.py @@ -12,6 +12,10 @@ import os import shepherd.scheduler +class PluginLoadError(Exception): + pass + + class Hook(): def __init__(self): self.attached_functions = [] @@ -242,17 +246,17 @@ def find_plugins(plugin_names, plugin_dir=None): # that can't be found except ModuleNotFoundError: try: - if (plugin_dir is not None) and (plugin_dir != ""): + if 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") + raise PluginLoadError("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) + raise PluginLoadError("Could not find plugin "+plugin_name) # Scan imported module for Plugin subclass attrs = [getattr(mod, name) for name in dir(mod)] @@ -262,8 +266,8 @@ def find_plugins(plugin_names, plugin_dir=None): 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") + raise PluginLoadError("Imported shepherd plugin modules must contain a " + "subclass of shepherd.plugin.Plugin, such as" + "shepherd.plugin.SimplePlugin") return plugin_classes