# Copyright 2014-2015 Canonical Limited. # # This file is part of charm-helpers. # # charm-helpers is free software: you can redistribute it and/or modify # it under the terms of the GNU Lesser General Public License version 3 as # published by the Free Software Foundation. # # charm-helpers is distributed in the hope that it will be useful, # but WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # GNU Lesser General Public License for more details. # # You should have received a copy of the GNU Lesser General Public License # along with charm-helpers. If not, see . import inspect import argparse import sys from six.moves import zip import charmhelpers.core.unitdata class OutputFormatter(object): def __init__(self, outfile=sys.stdout): self.formats = ( "raw", "json", "py", "yaml", "csv", "tab", ) self.outfile = outfile def add_arguments(self, argument_parser): formatgroup = argument_parser.add_mutually_exclusive_group() choices = self.supported_formats formatgroup.add_argument("--format", metavar='FMT', help="Select output format for returned data, " "where FMT is one of: {}".format(choices), choices=choices, default='raw') for fmt in self.formats: fmtfunc = getattr(self, fmt) formatgroup.add_argument("-{}".format(fmt[0]), "--{}".format(fmt), action='store_const', const=fmt, dest='format', help=fmtfunc.__doc__) @property def supported_formats(self): return self.formats def raw(self, output): """Output data as raw string (default)""" if isinstance(output, (list, tuple)): output = '\n'.join(map(str, output)) self.outfile.write(str(output)) def py(self, output): """Output data as a nicely-formatted python data structure""" import pprint pprint.pprint(output, stream=self.outfile) def json(self, output): """Output data in JSON format""" import json json.dump(output, self.outfile) def yaml(self, output): """Output data in YAML format""" import yaml yaml.safe_dump(output, self.outfile) def csv(self, output): """Output data as excel-compatible CSV""" import csv csvwriter = csv.writer(self.outfile) csvwriter.writerows(output) def tab(self, output): """Output data in excel-compatible tab-delimited format""" import csv csvwriter = csv.writer(self.outfile, dialect=csv.excel_tab) csvwriter.writerows(output) def format_output(self, output, fmt='raw'): fmtfunc = getattr(self, fmt) fmtfunc(output) class CommandLine(object): argument_parser = None subparsers = None formatter = None exit_code = 0 def __init__(self): if not self.argument_parser: self.argument_parser = argparse.ArgumentParser(description='Perform common charm tasks') if not self.formatter: self.formatter = OutputFormatter() self.formatter.add_arguments(self.argument_parser) if not self.subparsers: self.subparsers = self.argument_parser.add_subparsers(help='Commands') def subcommand(self, command_name=None): """ Decorate a function as a subcommand. Use its arguments as the command-line arguments""" def wrapper(decorated): cmd_name = command_name or decorated.__name__ subparser = self.subparsers.add_parser(cmd_name, description=decorated.__doc__) for args, kwargs in describe_arguments(decorated): subparser.add_argument(*args, **kwargs) subparser.set_defaults(func=decorated) return decorated return wrapper def test_command(self, decorated): """ Subcommand is a boolean test function, so bool return values should be converted to a 0/1 exit code. """ decorated._cli_test_command = True return decorated def no_output(self, decorated): """ Subcommand is not expected to return a value, so don't print a spurious None. """ decorated._cli_no_output = True return decorated def subcommand_builder(self, command_name, description=None): """ Decorate a function that builds a subcommand. Builders should accept a single argument (the subparser instance) and return the function to be run as the command.""" def wrapper(decorated): subparser = self.subparsers.add_parser(command_name) func = decorated(subparser) subparser.set_defaults(func=func) subparser.description = description or func.__doc__ return wrapper def run(self): "Run cli, processing arguments and executing subcommands." arguments = self.argument_parser.parse_args() argspec = inspect.getargspec(arguments.func) vargs = [] for arg in argspec.args: vargs.append(getattr(arguments, arg)) if argspec.varargs: vargs.extend(getattr(arguments, argspec.varargs)) output = arguments.func(*vargs) if getattr(arguments.func, '_cli_test_command', False): self.exit_code = 0 if output else 1 output = '' if getattr(arguments.func, '_cli_no_output', False): output = '' self.formatter.format_output(output, arguments.format) if charmhelpers.core.unitdata._KV: charmhelpers.core.unitdata._KV.flush() cmdline = CommandLine() def describe_arguments(func): """ Analyze a function's signature and return a data structure suitable for passing in as arguments to an argparse parser's add_argument() method.""" argspec = inspect.getargspec(func) # we should probably raise an exception somewhere if func includes **kwargs if argspec.defaults: positional_args = argspec.args[:-len(argspec.defaults)] keyword_names = argspec.args[-len(argspec.defaults):] for arg, default in zip(keyword_names, argspec.defaults): yield ('--{}'.format(arg),), {'default': default} else: positional_args = argspec.args for arg in positional_args: yield (arg,), {} if argspec.varargs: yield (argspec.varargs,), {'nargs': '*'}