Using variable arg names with argparse

问题:

I have an upstream system that invokes my program with varying arg names. Example:

foo --xyz1 10 --xyz2 25 --xyz3 31

I would like the result of argparsing to be a xyz = [10, 25, 31].

The names of my args have a common prefix, but unfortunately have to differ at least with a different numeric suffix, which also indicates order. I also don't have a fixed number of args.

Is there a way to model this with argparse? Either with what is available through some combination of built-in capabilities, or by overriding/pluging in some custom parser handling.

回答1:

I would suggest a bit of pre-processing to achieve this:

Code:

def get_xyz_cmd_line(xyz_cmd_line):
    # build a generator to iterate the cmd_line
    cmd_line_gen = iter(xyz_cmd_line)

    # we will separate the xyz's from everything else
    xyz = []
    remaining_cmd_line = []

    # go through the command line and extract the xyz's
    for opt in cmd_line_gen:
        if opt.startswith('--xyz'):
            # grab the opt and the arg for it
            xyz.append((opt, cmd_line_gen.next()))
        else:
            remaining_cmd_line.append(opt)

    # sort the xyz's and return all of them as -xyz # -xyz # ... 
    return list(it.chain(*[
        ('--xyz', x[1]) for x in sorted(xyz)])) + remaining_cmd_line 

To Test:

import argparse
import itertools as it

parser = argparse.ArgumentParser(description='Get my Option')
parser.add_argument('--an_opt', metavar='N', type=int,
                    help='An option')
parser.add_argument('--xyz', metavar='N', type=int, action='append',
                    help='An option')

cmd_line = "--an_opt 1 --xyz1 10 --xyz3 31 --xyz2 25 ".split()
args = parser.parse_args(get_xyz_cmd_line(cmd_line))
print(args)

Output:

Namespace(an_opt=1, xyz=[10, 25, 31])

To use:

Nominally instead of a fixed cmd_line as in the above example this would be called with something like:

args = parser.parse_args(get_xyz_cmd_line(sys.argv[1:]))

UPDATE: If you need --xyz=31 (ie = separator):

Then you will need to change:

# grab the opt and the arg for it
xyz.append((opt, cmd_line_gen.next()))

To:

if '=' in opt:
    xyz.append(tuple(opt.split('=', 1)))
else:
    # grab the opt and the arg for it
    xyz.append((opt, cmd_line_gen.next()))

回答2:

Here's what I did for reference (quick and dirty version), though I also like Stephen Rauch's answer (so I'll mark that as an answer -- esp. since I used internal implementation detail for my solution):

class CustomArgumentsParser(argparse.ArgumentParser):

  def _parse_optional(self, arg_string):
    suffix_index = arg_string.find(':')
    if suffix_index < 0:
      return super(CustomArgumentParser, self)._parse_optional(arg_string)

    original_arg_string = arg_string
    suffix = arg_string[suffix_index + 1:]
    arg_string = arg_string[0:suffix_index]

    option_tuple = super(CustomArgumentParser, self)._parse_optional(arg_string)
    if not option_tuple:
      return option_tuple

    action, option_string, explicit_arg = option_tuple
    if isinstance(action, BuildListAction):
      return action, suffix, explicit_arg
    else:
      self.exit(-1, message='Unknown argument %s' % original_arg_string)


class BuildListAction(argparse.Action):
  def __init__(self,
               option_strings,
               dest,
               nargs=None,
               const=None,
               default=None,
               type=None,
               choices=None,
               required=False,
               help=None,
               metavar=None):
    super(BuildListAction, self).__init__(
      option_strings=option_strings,
      dest=dest,
      nargs=nargs,
      const=const,
      default=default,
      type=type,
      choices=choices,
      required=required,
      help=help,
      metavar=metavar)

  def __call__(self, parser, namespace, values, option_string=None):
    index = int(option_string) - 1

    list = getattr(namespace, self.dest)
    if list is None:
      list = []
      setattr(namespace, self.dest, list)

    if index >= len(list):
      list.extend([self.default] * (index + 1 - len(list)))
    list[index] = values

Usage:

argparser = CustomArgumentsParser()
argparser.add_argument('--xyz', type=int, action=BuildListAction)

Note -- This supports args of the form --xyz:1, --xyz:2, ... which is slightly different than the original question.