argparse_oe.py 7.0 KB

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