1 files changed, 59 insertions, 0 deletions

diff --git a/scripts/lib/argparse_oe.py b/scripts/lib/argparse_oe.py
index fd866922bd..744cfe312f 100644
--- a/scripts/lib/argparse_oe.py
+++ b/scripts/lib/argparse_oe.py
@@ -1,5 +1,6 @@
 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"""
@@ -9,6 +10,10 @@ class ArgumentUsageError(Exception):
 
 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)
@@ -27,10 +32,26 @@ class ArgumentParser(argparse.ArgumentParser):
 
     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)
+
     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
@@ -64,3 +85,41 @@ class ArgumentSubParser(ArgumentParser):
             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)