aboutsummaryrefslogtreecommitdiffstats
path: root/scripts/lib/argparse_oe.py
blob: bf6eb17197b89f4502eb3523b577517aa886b5a1 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
import sys
import argparse
from collections import defaultdict, OrderedDict

class ArgumentUsageError(Exception):
    """Exception class you can raise (and catch) in order to show the help"""
    def __init__(self, message, subcommand=None):
        self.message = message
        self.subcommand = subcommand

class ArgumentParser(argparse.ArgumentParser):
    """Our own version of argparse's ArgumentParser"""
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('formatter_class', OeHelpFormatter)
        self._subparser_groups = OrderedDict()
        super(ArgumentParser, self).__init__(*args, **kwargs)
        self._positionals.title = 'arguments'
        self._optionals.title = 'options'

    def error(self, message):
        """error(message: string)

        Prints a help message incorporating the message to stderr and
        exits.
        """
        self._print_message('%s: error: %s\n' % (self.prog, message), sys.stderr)
        self.print_help(sys.stderr)
        sys.exit(2)

    def error_subcommand(self, message, subcommand):
        if subcommand:
            action = self._get_subparser_action()
            try:
                subparser = action._name_parser_map[subcommand]
            except KeyError:
                self.error('no subparser for name "%s"' % subcommand)
            else:
                subparser.error(message)

        self.error(message)

    def add_subparsers(self, *args, **kwargs):
        if 'dest' not in kwargs:
            kwargs['dest'] = '_subparser_name'

        ret = super(ArgumentParser, self).add_subparsers(*args, **kwargs)
        # Need a way of accessing the parent parser
        ret._parent_parser = self
        # Ensure our class gets instantiated
        ret._parser_class = ArgumentSubParser
        # Hacky way of adding a method to the subparsers object
        ret.add_subparser_group = self.add_subparser_group
        return ret

    def add_subparser_group(self, groupname, groupdesc, order=0):
        self._subparser_groups[groupname] = (groupdesc, order)

    def parse_args(self, args=None, namespace=None):
        """Parse arguments, using the correct subparser to show the error."""
        args, argv = self.parse_known_args(args, namespace)
        if argv:
            message = 'unrecognized arguments: %s' % ' '.join(argv)
            if self._subparsers:
                subparser = self._get_subparser(args)
                subparser.error(message)
            else:
                self.error(message)
            sys.exit(2)
        return args

    def _get_subparser(self, args):
        action = self._get_subparser_action()
        if action.dest == argparse.SUPPRESS:
            self.error('cannot get subparser, the subparser action dest is suppressed')

        name = getattr(args, action.dest)
        try:
            return action._name_parser_map[name]
        except KeyError:
            self.error('no subparser for name "%s"' % name)

    def _get_subparser_action(self):
        if not self._subparsers:
            self.error('cannot return the subparser action, no subparsers added')

        for action in self._subparsers._group_actions:
            if isinstance(action, argparse._SubParsersAction):
                return action


class ArgumentSubParser(ArgumentParser):
    def __init__(self, *args, **kwargs):
        if 'group' in kwargs:
            self._group = kwargs.pop('group')
        if 'order' in kwargs:
            self._order = kwargs.pop('order')
        super(ArgumentSubParser, self).__init__(*args, **kwargs)

    def parse_known_args(self, args=None, namespace=None):
        # This works around argparse not handling optional positional arguments being
        # intermixed with other options. A pretty horrible hack, but we're not left
        # with much choice given that the bug in argparse exists and it's difficult
        # to subclass.
        # Borrowed from http://stackoverflow.com/questions/20165843/argparse-how-to-handle-variable-number-of-arguments-nargs
        # with an extra workaround (in format_help() below) for the positional
        # arguments disappearing from the --help output, as well as structural tweaks.
        # Originally simplified from http://bugs.python.org/file30204/test_intermixed.py
        positionals = self._get_positional_actions()
        for action in positionals:
            # deactivate positionals
            action.save_nargs = action.nargs
            action.nargs = 0

        namespace, remaining_args = super(ArgumentSubParser, self).parse_known_args(args, namespace)
        for action in positionals:
            # remove the empty positional values from namespace
            if hasattr(namespace, action.dest):
                delattr(namespace, action.dest)
        for action in positionals:
            action.nargs = action.save_nargs
        # parse positionals
        namespace, extras = super(ArgumentSubParser, self).parse_known_args(remaining_args, namespace)
        return namespace, extras

    def format_help(self):
        # Quick, restore the positionals!
        positionals = self._get_positional_actions()
        for action in positionals:
            if hasattr(action, 'save_nargs'):
                action.nargs = action.save_nargs
        return super(ArgumentParser, self).format_help()


class OeHelpFormatter(argparse.HelpFormatter):
    def _format_action(self, action):
        if hasattr(action, '_get_subactions'):
            # subcommands list
            groupmap = defaultdict(list)
            ordermap = {}
            subparser_groups = action._parent_parser._subparser_groups
            groups = sorted(subparser_groups.keys(), key=lambda item: subparser_groups[item][1], reverse=True)
            for subaction in self._iter_indented_subactions(action):
                parser = action._name_parser_map[subaction.dest]
                group = getattr(parser, '_group', None)
                groupmap[group].append(subaction)
                if group not in groups:
                    groups.append(group)
                order = getattr(parser, '_order', 0)
                ordermap[subaction.dest] = order

            lines = []
            if len(groupmap) > 1:
                groupindent = '  '
            else:
                groupindent = ''
            for group in groups:
                subactions = groupmap[group]
                if not subactions:
                    continue
                if groupindent:
                    if not group:
                        group = 'other'
                    groupdesc = subparser_groups.get(group, (group, 0))[0]
                    lines.append('  %s:' % groupdesc)
                for subaction in sorted(subactions, key=lambda item: ordermap[item.dest], reverse=True):
                    lines.append('%s%s' % (groupindent, self._format_action(subaction).rstrip()))
            return '\n'.join(lines)
        else:
            return super(OeHelpFormatter, self)._format_action(action)