Add AliasedSubParsersAction

qubes/tools/__init__.py

@@ -291,6 +291,45 @@ class QubesArgumentParser(argparse.ArgumentParser):
         print(*args, file=sys.stderr, **kwargs)
 
 
+class AliasedSubParsersAction(argparse._SubParsersAction):
+    # source https://gist.github.com/sampsyo/471779
+    # pylint: disable=protected-access,too-few-public-methods
+    class _AliasedPseudoAction(argparse.Action):
+        # pylint: disable=redefined-builtin
+        def __init__(self, name, aliases, help):
+            dest = name
+            if aliases:
+                dest += ' (%s)' % ','.join(aliases)
+            sup = super(AliasedSubParsersAction._AliasedPseudoAction, self)
+            sup.__init__(option_strings=[], dest=dest, help=help)
+
+        def __call__(self, **kwargs):
+            super(AliasedSubParsersAction._AliasedPseudoAction, self).__call__(
+                **kwargs)
+
+    def add_parser(self, name, **kwargs):
+        if 'aliases' in kwargs:
+            aliases = kwargs['aliases']
+            del kwargs['aliases']
+        else:
+            aliases = []
+
+        local_parser = super(AliasedSubParsersAction, self).add_parser(
+            name, **kwargs)
+
+        # Make the aliases work.
+        for alias in aliases:
+            self._name_parser_map[alias] = local_parser
+        # Make the help text reflect them, first removing old help entry.
+        if 'help' in kwargs:
+            self._choices_actions.pop()
+            pseudo_action = self._AliasedPseudoAction(name, aliases,
+                                                      kwargs.pop('help'))
+            self._choices_actions.append(pseudo_action)
+
+        return local_parser
+
+
 def get_parser_for_command(command):
     '''Get parser for given qvm-tool.