summaryrefslogtreecommitdiff
path: root/scripts/lib/argparse_oe.py
blob: bf3ebaddfd9b244acfa319c1bcf8e8896ff34704 (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
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)

    def error(self, message):
        sys.stderr.write('ERROR: %s\n' % message)
        self.print_help()
        sys.exit(2)

    def error_subcommand(self, message, subcommand):
        if subcommand:
            for action in self._actions:
                if isinstance(action, argparse._SubParsersAction):
                    for choice, subparser in action.choices.items():
                        if choice == subcommand:
                            subparser.error(message)
                            return
        self.error(message)

    def add_subparsers(self, *args, **kwargs):
        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)


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)
        for agroup in self._action_groups:
            if agroup.title == 'optional arguments':
                agroup.title = 'options'
                break

    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)