如何使用argparse解析1-3指定的范围并返回列表

Question

嗨，简单来说，我想在命令行中添加一个 --cid 参数。当我输入 --cid 1-3 7 22 时，我希望解析出来的参数是一个列表 [1,2,3,7,22]。

我现在的尝试是：

from argparse import ArgumentParser, ArgumentTypeError
import re

def parseNumList(string):
    print 'in string:',string
    if '-' in string:
        m = re.match(r'(\d+)(?:-(\d+))?$', string)
        # ^ (or use .split('-'). anyway you like.)
        if not m:
            raise ArgumentTypeError("'" + string + "' is not a range of number. Expected forms like '0-5' or '2'.")
        start = m.group(1)
        end = m.group(2) or start
        return list(range(int(start,10), int(end,10)+1))
    else:
        return string

parser = ArgumentParser()
parser.add_argument('--cid', type=parseNumList,nargs='*')

args = parser.parse_args()
print(args)

但是结果是这样，我不想要嵌套列表，我该怎么做呢？（我知道我可以手动解析 args.cid 并重新赋值，但我能不能在 argparse 里直接完成这一切？使用自定义动作或自定义类型？）

> python testArgs.py --cid 1-3 7 22
in string: 1-3
in string: 7
in string: 22
Namespace(cid=[[1, 2, 3], '7', '22'])

编辑：我用自定义动作搞定了，但我觉得这不是最好的方法。有什么建议吗？

from argparse import ArgumentParser, ArgumentTypeError,Action
import re

class CustomRangeAction(Action):
    def __call__(self, parser, namespace, values, option_string=None):
        print 'running action for ' + repr(values)

        flatten = []

        for v in values:
            try:
                m = re.match(r'(\d+)(?:-(\d+))?$', v)  
                if m:
                    start = m.group(1)
                    end = m.group(2) or start
                    flatten.extend(list(range(int(start,10), int(end,10)+1)))
                else:
                    flatten.append(int(v))
            except:
                continue

        flatten = sorted(list(set(flatten)))
        flatten.reverse()

        setattr(namespace, self.dest, flatten)


parser = ArgumentParser()
parser.add_argument('--cid', action=CustomRangeAction,nargs='*')

args = parser.parse_args()
print(args)

Answer 1

这个回答“正确的做法是...”说得不错，但我需要修改一下InflateRange类的代码，才能让它真正工作。

另外，为了把错误控制在argparse内部，抛出的异常需要是ArgumentError类型的（可以参考这个回答）。

自定义操作的代码看起来是这样的：

class InflateRange(argparse.Action):
    def __call__(self, parser, namespace, values, option_string=None):
        lst = []

        limits = values.split('-')
        if len(limits) > 2:
            raise argparse.ArgumentError(self, "%r is not a range of numbers." % part )
        start = limits[0]
        if start == '': start = '1'
        end = limits[-1]
        try:
            lst.extend(list(range(int(start), int(end)+1)))
        except ValueError:
            raise argparse.ArgumentError(self, "%r is not a range of numbers." % part )

        setattr(namespace, self.dest, lst)

Answer 2

你可以用下面的代码来实现你的 parseNumList 方法。我假设 input_string 是这个方法的输入。

>>> input_string = '1-3 7 22'
>>> input_list = input_string.split()
>>> alist = []
>>> for item in input_list:
...     if '-' in item:
...             a, b = item.split('-')
...             alist.extend(range(int(a), int(b) + 1))
...     else:
...             alist.append(int(item))
...
>>> alist
[1, 2, 3, 7, 22]

你可以在合适的地方添加错误处理和边界检查。

Answer 3

正确的做法是使用 action= 而不是 type= 在 add_argument 中来定义特殊的行为。当每个参数都是一种特殊类型时（比如，每个参数是一个表示数字的十六进制字符串），你应该使用 type。在这种情况下，你的一些参数需要被扩展成多个。

为了定义一个动作，你需要定义一个新的类（从 argparse.Action 继承），并实现 __call__(self, parser, namespace, values, option_string=None) 方法：

from argparse import ArgumentParser, ArgumentTypeError, Action
import re

class InflateRange(Action):
    def __call__(self, parser, namespace, values, option_string=None):
        print('%r %r %r' % (namespace, values, option_string))
        lst = []
        for string in values:
            print 'in string:',string
            if '-' in string:
                m = re.match(r'(\d+)(?:-(\d+))?$', string)
                # ^ (or use .split('-'). anyway you like.)
                if not m:
                    raise ArgumentTypeError("'" + string + "' is not a range of number. Expected forms like '0-5' or '2'.")
                start = m.group(1)
                end = m.group(2) or start
                lst.extend(list(range(int(start,10), int(end,10)+1)))
            else:
                 lst.append(int(string))
            setattr(namespace, self.dest, lst)

parser = ArgumentParser()
parser.add_argument('--cid', action=InflateRange, nargs='*')

args = parser.parse_args()
print(args)

这样你就能得到：

> python test.py --cid 1-3 7 22
Namespace(cid=None) ['1-3', '7', '22'] '--cid'
in string: 1-3
in string: 7
in string: 22
Namespace(cid=[1, 2, 3, 7, 22])

想了解更多信息，可以查看 文档