~james-page/charms/trusty/heat/lp1531102-trunk

« back to all changes in this revision

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

  • Committer: james.page at ubuntu
  • Date: 2015-08-10 16:35:15 UTC
  • Revision ID: james.page@ubuntu.com-20150810163515-a3d9pwggu4aktvvf
Tags: 15.07
[gnuoy] 15.07 Charm release

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
# Copyright 2014-2015 Canonical Limited.
 
2
#
 
3
# This file is part of charm-helpers.
 
4
#
 
5
# charm-helpers is free software: you can redistribute it and/or modify
 
6
# it under the terms of the GNU Lesser General Public License version 3 as
 
7
# published by the Free Software Foundation.
 
8
#
 
9
# charm-helpers is distributed in the hope that it will be useful,
 
10
# but WITHOUT ANY WARRANTY; without even the implied warranty of
 
11
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 
12
# GNU Lesser General Public License for more details.
 
13
#
 
14
# You should have received a copy of the GNU Lesser General Public License
 
15
# along with charm-helpers.  If not, see <http://www.gnu.org/licenses/>.
 
16
 
 
17
import inspect
 
18
import argparse
 
19
import sys
 
20
 
 
21
from six.moves import zip
 
22
 
 
23
from charmhelpers.core import unitdata
 
24
 
 
25
 
 
26
class OutputFormatter(object):
 
27
    def __init__(self, outfile=sys.stdout):
 
28
        self.formats = (
 
29
            "raw",
 
30
            "json",
 
31
            "py",
 
32
            "yaml",
 
33
            "csv",
 
34
            "tab",
 
35
        )
 
36
        self.outfile = outfile
 
37
 
 
38
    def add_arguments(self, argument_parser):
 
39
        formatgroup = argument_parser.add_mutually_exclusive_group()
 
40
        choices = self.supported_formats
 
41
        formatgroup.add_argument("--format", metavar='FMT',
 
42
                                 help="Select output format for returned data, "
 
43
                                      "where FMT is one of: {}".format(choices),
 
44
                                 choices=choices, default='raw')
 
45
        for fmt in self.formats:
 
46
            fmtfunc = getattr(self, fmt)
 
47
            formatgroup.add_argument("-{}".format(fmt[0]),
 
48
                                     "--{}".format(fmt), action='store_const',
 
49
                                     const=fmt, dest='format',
 
50
                                     help=fmtfunc.__doc__)
 
51
 
 
52
    @property
 
53
    def supported_formats(self):
 
54
        return self.formats
 
55
 
 
56
    def raw(self, output):
 
57
        """Output data as raw string (default)"""
 
58
        if isinstance(output, (list, tuple)):
 
59
            output = '\n'.join(map(str, output))
 
60
        self.outfile.write(str(output))
 
61
 
 
62
    def py(self, output):
 
63
        """Output data as a nicely-formatted python data structure"""
 
64
        import pprint
 
65
        pprint.pprint(output, stream=self.outfile)
 
66
 
 
67
    def json(self, output):
 
68
        """Output data in JSON format"""
 
69
        import json
 
70
        json.dump(output, self.outfile)
 
71
 
 
72
    def yaml(self, output):
 
73
        """Output data in YAML format"""
 
74
        import yaml
 
75
        yaml.safe_dump(output, self.outfile)
 
76
 
 
77
    def csv(self, output):
 
78
        """Output data as excel-compatible CSV"""
 
79
        import csv
 
80
        csvwriter = csv.writer(self.outfile)
 
81
        csvwriter.writerows(output)
 
82
 
 
83
    def tab(self, output):
 
84
        """Output data in excel-compatible tab-delimited format"""
 
85
        import csv
 
86
        csvwriter = csv.writer(self.outfile, dialect=csv.excel_tab)
 
87
        csvwriter.writerows(output)
 
88
 
 
89
    def format_output(self, output, fmt='raw'):
 
90
        fmtfunc = getattr(self, fmt)
 
91
        fmtfunc(output)
 
92
 
 
93
 
 
94
class CommandLine(object):
 
95
    argument_parser = None
 
96
    subparsers = None
 
97
    formatter = None
 
98
    exit_code = 0
 
99
 
 
100
    def __init__(self):
 
101
        if not self.argument_parser:
 
102
            self.argument_parser = argparse.ArgumentParser(description='Perform common charm tasks')
 
103
        if not self.formatter:
 
104
            self.formatter = OutputFormatter()
 
105
            self.formatter.add_arguments(self.argument_parser)
 
106
        if not self.subparsers:
 
107
            self.subparsers = self.argument_parser.add_subparsers(help='Commands')
 
108
 
 
109
    def subcommand(self, command_name=None):
 
110
        """
 
111
        Decorate a function as a subcommand. Use its arguments as the
 
112
        command-line arguments"""
 
113
        def wrapper(decorated):
 
114
            cmd_name = command_name or decorated.__name__
 
115
            subparser = self.subparsers.add_parser(cmd_name,
 
116
                                                   description=decorated.__doc__)
 
117
            for args, kwargs in describe_arguments(decorated):
 
118
                subparser.add_argument(*args, **kwargs)
 
119
            subparser.set_defaults(func=decorated)
 
120
            return decorated
 
121
        return wrapper
 
122
 
 
123
    def test_command(self, decorated):
 
124
        """
 
125
        Subcommand is a boolean test function, so bool return values should be
 
126
        converted to a 0/1 exit code.
 
127
        """
 
128
        decorated._cli_test_command = True
 
129
        return decorated
 
130
 
 
131
    def no_output(self, decorated):
 
132
        """
 
133
        Subcommand is not expected to return a value, so don't print a spurious None.
 
134
        """
 
135
        decorated._cli_no_output = True
 
136
        return decorated
 
137
 
 
138
    def subcommand_builder(self, command_name, description=None):
 
139
        """
 
140
        Decorate a function that builds a subcommand. Builders should accept a
 
141
        single argument (the subparser instance) and return the function to be
 
142
        run as the command."""
 
143
        def wrapper(decorated):
 
144
            subparser = self.subparsers.add_parser(command_name)
 
145
            func = decorated(subparser)
 
146
            subparser.set_defaults(func=func)
 
147
            subparser.description = description or func.__doc__
 
148
        return wrapper
 
149
 
 
150
    def run(self):
 
151
        "Run cli, processing arguments and executing subcommands."
 
152
        arguments = self.argument_parser.parse_args()
 
153
        argspec = inspect.getargspec(arguments.func)
 
154
        vargs = []
 
155
        kwargs = {}
 
156
        for arg in argspec.args:
 
157
            vargs.append(getattr(arguments, arg))
 
158
        if argspec.varargs:
 
159
            vargs.extend(getattr(arguments, argspec.varargs))
 
160
        if argspec.keywords:
 
161
            for kwarg in argspec.keywords.items():
 
162
                kwargs[kwarg] = getattr(arguments, kwarg)
 
163
        output = arguments.func(*vargs, **kwargs)
 
164
        if getattr(arguments.func, '_cli_test_command', False):
 
165
            self.exit_code = 0 if output else 1
 
166
            output = ''
 
167
        if getattr(arguments.func, '_cli_no_output', False):
 
168
            output = ''
 
169
        self.formatter.format_output(output, arguments.format)
 
170
        if unitdata._KV:
 
171
            unitdata._KV.flush()
 
172
 
 
173
 
 
174
cmdline = CommandLine()
 
175
 
 
176
 
 
177
def describe_arguments(func):
 
178
    """
 
179
    Analyze a function's signature and return a data structure suitable for
 
180
    passing in as arguments to an argparse parser's add_argument() method."""
 
181
 
 
182
    argspec = inspect.getargspec(func)
 
183
    # we should probably raise an exception somewhere if func includes **kwargs
 
184
    if argspec.defaults:
 
185
        positional_args = argspec.args[:-len(argspec.defaults)]
 
186
        keyword_names = argspec.args[-len(argspec.defaults):]
 
187
        for arg, default in zip(keyword_names, argspec.defaults):
 
188
            yield ('--{}'.format(arg),), {'default': default}
 
189
    else:
 
190
        positional_args = argspec.args
 
191
 
 
192
    for arg in positional_args:
 
193
        yield (arg,), {}
 
194
    if argspec.varargs:
 
195
        yield (argspec.varargs,), {'nargs': '*'}