blob: bf6eb17197b89f4502eb3523b577517aa886b5a1 [file] [log] [blame]
Patrick Williamsd8c66bc2016-06-20 12:57:21 -05001import sys
2import argparse
3from collections import defaultdict, OrderedDict
4
5class ArgumentUsageError(Exception):
6 """Exception class you can raise (and catch) in order to show the help"""
7 def __init__(self, message, subcommand=None):
8 self.message = message
9 self.subcommand = subcommand
10
11class ArgumentParser(argparse.ArgumentParser):
12 """Our own version of argparse's ArgumentParser"""
13 def __init__(self, *args, **kwargs):
14 kwargs.setdefault('formatter_class', OeHelpFormatter)
15 self._subparser_groups = OrderedDict()
16 super(ArgumentParser, self).__init__(*args, **kwargs)
Patrick Williamsc0f7c042017-02-23 20:41:17 -060017 self._positionals.title = 'arguments'
18 self._optionals.title = 'options'
Patrick Williamsd8c66bc2016-06-20 12:57:21 -050019
20 def error(self, message):
Patrick Williamsc0f7c042017-02-23 20:41:17 -060021 """error(message: string)
22
23 Prints a help message incorporating the message to stderr and
24 exits.
25 """
26 self._print_message('%s: error: %s\n' % (self.prog, message), sys.stderr)
27 self.print_help(sys.stderr)
Patrick Williamsd8c66bc2016-06-20 12:57:21 -050028 sys.exit(2)
29
30 def error_subcommand(self, message, subcommand):
31 if subcommand:
Patrick Williamsc0f7c042017-02-23 20:41:17 -060032 action = self._get_subparser_action()
33 try:
34 subparser = action._name_parser_map[subcommand]
35 except KeyError:
36 self.error('no subparser for name "%s"' % subcommand)
37 else:
38 subparser.error(message)
39
Patrick Williamsd8c66bc2016-06-20 12:57:21 -050040 self.error(message)
41
42 def add_subparsers(self, *args, **kwargs):
Patrick Williamsc0f7c042017-02-23 20:41:17 -060043 if 'dest' not in kwargs:
44 kwargs['dest'] = '_subparser_name'
45
Patrick Williamsd8c66bc2016-06-20 12:57:21 -050046 ret = super(ArgumentParser, self).add_subparsers(*args, **kwargs)
47 # Need a way of accessing the parent parser
48 ret._parent_parser = self
49 # Ensure our class gets instantiated
50 ret._parser_class = ArgumentSubParser
51 # Hacky way of adding a method to the subparsers object
52 ret.add_subparser_group = self.add_subparser_group
53 return ret
54
55 def add_subparser_group(self, groupname, groupdesc, order=0):
56 self._subparser_groups[groupname] = (groupdesc, order)
57
Patrick Williamsc0f7c042017-02-23 20:41:17 -060058 def parse_args(self, args=None, namespace=None):
59 """Parse arguments, using the correct subparser to show the error."""
60 args, argv = self.parse_known_args(args, namespace)
61 if argv:
62 message = 'unrecognized arguments: %s' % ' '.join(argv)
63 if self._subparsers:
64 subparser = self._get_subparser(args)
65 subparser.error(message)
66 else:
67 self.error(message)
68 sys.exit(2)
69 return args
70
71 def _get_subparser(self, args):
72 action = self._get_subparser_action()
73 if action.dest == argparse.SUPPRESS:
74 self.error('cannot get subparser, the subparser action dest is suppressed')
75
76 name = getattr(args, action.dest)
77 try:
78 return action._name_parser_map[name]
79 except KeyError:
80 self.error('no subparser for name "%s"' % name)
81
82 def _get_subparser_action(self):
83 if not self._subparsers:
84 self.error('cannot return the subparser action, no subparsers added')
85
86 for action in self._subparsers._group_actions:
87 if isinstance(action, argparse._SubParsersAction):
88 return action
89
Patrick Williamsd8c66bc2016-06-20 12:57:21 -050090
91class ArgumentSubParser(ArgumentParser):
92 def __init__(self, *args, **kwargs):
93 if 'group' in kwargs:
94 self._group = kwargs.pop('group')
95 if 'order' in kwargs:
96 self._order = kwargs.pop('order')
97 super(ArgumentSubParser, self).__init__(*args, **kwargs)
Patrick Williamsd8c66bc2016-06-20 12:57:21 -050098
99 def parse_known_args(self, args=None, namespace=None):
100 # This works around argparse not handling optional positional arguments being
101 # intermixed with other options. A pretty horrible hack, but we're not left
102 # with much choice given that the bug in argparse exists and it's difficult
103 # to subclass.
104 # Borrowed from http://stackoverflow.com/questions/20165843/argparse-how-to-handle-variable-number-of-arguments-nargs
105 # with an extra workaround (in format_help() below) for the positional
106 # arguments disappearing from the --help output, as well as structural tweaks.
107 # Originally simplified from http://bugs.python.org/file30204/test_intermixed.py
108 positionals = self._get_positional_actions()
109 for action in positionals:
110 # deactivate positionals
111 action.save_nargs = action.nargs
112 action.nargs = 0
113
114 namespace, remaining_args = super(ArgumentSubParser, self).parse_known_args(args, namespace)
115 for action in positionals:
116 # remove the empty positional values from namespace
117 if hasattr(namespace, action.dest):
118 delattr(namespace, action.dest)
119 for action in positionals:
120 action.nargs = action.save_nargs
121 # parse positionals
122 namespace, extras = super(ArgumentSubParser, self).parse_known_args(remaining_args, namespace)
123 return namespace, extras
124
125 def format_help(self):
126 # Quick, restore the positionals!
127 positionals = self._get_positional_actions()
128 for action in positionals:
129 if hasattr(action, 'save_nargs'):
130 action.nargs = action.save_nargs
131 return super(ArgumentParser, self).format_help()
132
133
134class OeHelpFormatter(argparse.HelpFormatter):
135 def _format_action(self, action):
136 if hasattr(action, '_get_subactions'):
137 # subcommands list
138 groupmap = defaultdict(list)
139 ordermap = {}
140 subparser_groups = action._parent_parser._subparser_groups
141 groups = sorted(subparser_groups.keys(), key=lambda item: subparser_groups[item][1], reverse=True)
142 for subaction in self._iter_indented_subactions(action):
143 parser = action._name_parser_map[subaction.dest]
144 group = getattr(parser, '_group', None)
145 groupmap[group].append(subaction)
146 if group not in groups:
147 groups.append(group)
148 order = getattr(parser, '_order', 0)
149 ordermap[subaction.dest] = order
150
151 lines = []
152 if len(groupmap) > 1:
153 groupindent = ' '
154 else:
155 groupindent = ''
156 for group in groups:
157 subactions = groupmap[group]
158 if not subactions:
159 continue
160 if groupindent:
161 if not group:
162 group = 'other'
163 groupdesc = subparser_groups.get(group, (group, 0))[0]
164 lines.append(' %s:' % groupdesc)
165 for subaction in sorted(subactions, key=lambda item: ordermap[item.dest], reverse=True):
166 lines.append('%s%s' % (groupindent, self._format_action(subaction).rstrip()))
167 return '\n'.join(lines)
168 else:
169 return super(OeHelpFormatter, self)._format_action(action)