Python argparse второго выбора

Question

С помощью Python argparse модуля я в настоящее время это:

usage: prog [-h] (-a {opt1,opt2} | -b {opt3,opt4}) arg.

Однако, я хотел бы достичь следующего (второй набор вариантов для варианта -a):

usage: prog [-h] (-a {opt1,opt2} {opt5,opt6} | -b {opt3,opt4}) arg.

Я не знаю, как это сделать. Любая помощь?

Answer 1

Это работает:

import argparse 

class MultichoiceArgumentParser(argparse.ArgumentParser): 
    def _get_values(self, action, arg_strings): 
     if isinstance(action.nargs, int): 
      value = [self._get_value(action, v) for v in arg_strings] 
      for i, v in enumerate(value): 
       self._check_value(action, v, arg_num=i) 
      return value 
     else: 
      return super()._get_values(action, arg_strings) 

    def _check_value(self, action, value, arg_num=None): 
     # converted value must be one of the choices (if specified) 
     if action.choices is not None: 
      choices = action.choices 
      if isinstance(action, MultichoiceAction): 
       choices = choices[arg_num] 
      if value not in choices: 
       args = {'value': value, 
         'choices': ', '.join(map(repr, choices))} 
       msg = argparse._('invalid choice: %(value)r (choose from %(choices)s)') 
       raise argparse.ArgumentError(action, msg % args) 

class MultichoiceAction(argparse._StoreAction): 
    pass 

class MultichoiceFormatter(argparse.HelpFormatter): 
    def _metavar_formatter(self, action, default_metavar): 
     if action.choices is not None and isinstance(action, MultichoiceAction): 
      result = [] 
      for choices in action.choices: 
       choice_strs = [str(choice) for choice in choices] 
       result.append('{%s}' % ','.join(choice_strs)) 
      result = tuple(result) 
      return lambda x: result 
     else: 
      return super()._metavar_formatter(action, default_metavar) 

Чтобы использовать его:

parser = MultichoiceArgumentParser(formatter_class=MultichoiceFormatter) 
parser.add_argument("-a", choices=["ABC", "XYZ"], nargs=2, action=MultichoiceAction) 
print(parser.parse_args()) 

При использовании:

$ python d2.py -h 
usage: d2.py [-h] [-a {A,B,C} {X,Y,Z}] 

optional arguments: 
    -h, --help   show this help message and exit 
    -a {A,B,C} {X,Y,Z} 

$ python d2.py -a A X 
Namespace(a=['A', 'X']) 

$ python d2.py -a C Y 
Namespace(a=['C', 'Y']) 

$ python d2.py -a B Z 
Namespace(a=['B', 'Z']) 

$ python d2.py -a B B 
usage: d2.py [-h] [-a {A,B,C} {X,Y,Z}] 
d2.py: error: argument -a: invalid choice: 'B' (choose from 'X', 'Y', 'Z') 

$ python d2.py -a X X 
usage: d2.py [-h] [-a {A,B,C} {X,Y,Z}] 
d2.py: error: argument -a: invalid choice: 'X' (choose from 'A', 'B', 'C') 

$ python d2.py -a a b 
usage: d2.py [-h] [-a {A,B,C} {X,Y,Z}] 
d2.py: error: argument -a: invalid choice: 'a' (choose from 'A', 'B', 'C') 

Как показано выше, вы просто должны использовать MultichoiceArgumentParser с MultichoiceFormatter как аргумент formatter_class. Чтобы сделать аргумент с несколькими различными аргументами выбора, choice должен быть последовательностью последовательностей, такой же длины, как nargs, а action должен быть MultichoiceAction. Это действует так же, как store, поэтому для других видов поведения потребуется больше подкласса.