~tribaal/charms/precise/storage/refactor-mount-volume

« back to all changes in this revision

Viewing changes to hooks/charmhelpers/cli/__init__.py

  • Committer: David Britton
  • Date: 2014-02-05 20:46:42 UTC
  • mfrom: (26.1.89 storage)
  • Revision ID: dpb@canonical.com-20140205204642-qfwv0x6314bulcx7
merging chad's python/nfs additions

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
import inspect
 
2
import itertools
 
3
import argparse
 
4
import sys
 
5
 
 
6
 
 
7
class OutputFormatter(object):
 
8
    def __init__(self, outfile=sys.stdout):
 
9
        self.formats = (
 
10
            "raw",
 
11
            "json",
 
12
            "py",
 
13
            "yaml",
 
14
            "csv",
 
15
            "tab",
 
16
        )
 
17
        self.outfile = outfile
 
18
 
 
19
    def add_arguments(self, argument_parser):
 
20
        formatgroup = argument_parser.add_mutually_exclusive_group()
 
21
        choices = self.supported_formats
 
22
        formatgroup.add_argument("--format", metavar='FMT',
 
23
                                 help="Select output format for returned data, "
 
24
                                      "where FMT is one of: {}".format(choices),
 
25
                                 choices=choices, default='raw')
 
26
        for fmt in self.formats:
 
27
            fmtfunc = getattr(self, fmt)
 
28
            formatgroup.add_argument("-{}".format(fmt[0]),
 
29
                                     "--{}".format(fmt), action='store_const',
 
30
                                     const=fmt, dest='format',
 
31
                                     help=fmtfunc.__doc__)
 
32
 
 
33
    @property
 
34
    def supported_formats(self):
 
35
        return self.formats
 
36
 
 
37
    def raw(self, output):
 
38
        """Output data as raw string (default)"""
 
39
        self.outfile.write(str(output))
 
40
 
 
41
    def py(self, output):
 
42
        """Output data as a nicely-formatted python data structure"""
 
43
        import pprint
 
44
        pprint.pprint(output, stream=self.outfile)
 
45
 
 
46
    def json(self, output):
 
47
        """Output data in JSON format"""
 
48
        import json
 
49
        json.dump(output, self.outfile)
 
50
 
 
51
    def yaml(self, output):
 
52
        """Output data in YAML format"""
 
53
        import yaml
 
54
        yaml.safe_dump(output, self.outfile)
 
55
 
 
56
    def csv(self, output):
 
57
        """Output data as excel-compatible CSV"""
 
58
        import csv
 
59
        csvwriter = csv.writer(self.outfile)
 
60
        csvwriter.writerows(output)
 
61
 
 
62
    def tab(self, output):
 
63
        """Output data in excel-compatible tab-delimited format"""
 
64
        import csv
 
65
        csvwriter = csv.writer(self.outfile, dialect=csv.excel_tab)
 
66
        csvwriter.writerows(output)
 
67
 
 
68
    def format_output(self, output, fmt='raw'):
 
69
        fmtfunc = getattr(self, fmt)
 
70
        fmtfunc(output)
 
71
 
 
72
 
 
73
class CommandLine(object):
 
74
    argument_parser = None
 
75
    subparsers = None
 
76
    formatter = None
 
77
 
 
78
    def __init__(self):
 
79
        if not self.argument_parser:
 
80
            self.argument_parser = argparse.ArgumentParser(description='Perform common charm tasks')
 
81
        if not self.formatter:
 
82
            self.formatter = OutputFormatter()
 
83
            self.formatter.add_arguments(self.argument_parser)
 
84
        if not self.subparsers:
 
85
            self.subparsers = self.argument_parser.add_subparsers(help='Commands')
 
86
 
 
87
    def subcommand(self, command_name=None):
 
88
        """
 
89
        Decorate a function as a subcommand. Use its arguments as the
 
90
        command-line arguments"""
 
91
        def wrapper(decorated):
 
92
            cmd_name = command_name or decorated.__name__
 
93
            subparser = self.subparsers.add_parser(cmd_name,
 
94
                                                   description=decorated.__doc__)
 
95
            for args, kwargs in describe_arguments(decorated):
 
96
                subparser.add_argument(*args, **kwargs)
 
97
            subparser.set_defaults(func=decorated)
 
98
            return decorated
 
99
        return wrapper
 
100
 
 
101
    def subcommand_builder(self, command_name, description=None):
 
102
        """
 
103
        Decorate a function that builds a subcommand. Builders should accept a
 
104
        single argument (the subparser instance) and return the function to be
 
105
        run as the command."""
 
106
        def wrapper(decorated):
 
107
            subparser = self.subparsers.add_parser(command_name)
 
108
            func = decorated(subparser)
 
109
            subparser.set_defaults(func=func)
 
110
            subparser.description = description or func.__doc__
 
111
        return wrapper
 
112
 
 
113
    def run(self):
 
114
        "Run cli, processing arguments and executing subcommands."
 
115
        arguments = self.argument_parser.parse_args()
 
116
        argspec = inspect.getargspec(arguments.func)
 
117
        vargs = []
 
118
        kwargs = {}
 
119
        if argspec.varargs:
 
120
            vargs = getattr(arguments, argspec.varargs)
 
121
        for arg in argspec.args:
 
122
            kwargs[arg] = getattr(arguments, arg)
 
123
        self.formatter.format_output(arguments.func(*vargs, **kwargs), arguments.format)
 
124
 
 
125
 
 
126
cmdline = CommandLine()
 
127
 
 
128
 
 
129
def describe_arguments(func):
 
130
    """
 
131
    Analyze a function's signature and return a data structure suitable for
 
132
    passing in as arguments to an argparse parser's add_argument() method."""
 
133
 
 
134
    argspec = inspect.getargspec(func)
 
135
    # we should probably raise an exception somewhere if func includes **kwargs
 
136
    if argspec.defaults:
 
137
        positional_args = argspec.args[:-len(argspec.defaults)]
 
138
        keyword_names = argspec.args[-len(argspec.defaults):]
 
139
        for arg, default in itertools.izip(keyword_names, argspec.defaults):
 
140
            yield ('--{}'.format(arg),), {'default': default}
 
141
    else:
 
142
        positional_args = argspec.args
 
143
 
 
144
    for arg in positional_args:
 
145
        yield (arg,), {}
 
146
    if argspec.varargs:
 
147
        yield (argspec.varargs,), {'nargs': '*'}