python argparse 文件扩展名检查

Question

可以用 argparse 来验证命令行参数中的文件名扩展名吗？

比如说，我有一个从命令行运行的 Python 脚本：

$ script.py file.csv
$ script.py file.tab
$ script.py file.txt

我希望 argparse 能接受前两个文件名的命令行选项，但拒绝第三个。

我知道可以这样做：

parser = argparse.ArgumentParser()
parser.add_argument("fn", choices=["csv","tab"])
args = parser.parse_args()

来指定两个有效的命令行选项。

我想要的是这样的：

parser.add_argument("fn", choices=["*.csv","*.tab"])

来指定两个有效的文件扩展名作为命令行选项。不幸的是，这样做不行——有没有办法用 argparse 实现这个功能？

Answer 1

不可以。你可以给选择参数提供一个容器对象，或者任何支持“in”操作符的东西。想了解更多可以去看看这个链接。

不过，你总是可以自己检查一下，然后给用户反馈。

Answer 2

定义一个自定义函数，这个函数接收一个字符串形式的名字 - 把文件的扩展名分开以便进行比较，如果这个名字是可以接受的，就直接返回这个字符串；如果不可以，就抛出一个argparse所期望的异常：

def valid_file(param):
    base, ext = os.path.splitext(param)
    if ext.lower() not in ('.csv', '.tab'):
        raise argparse.ArgumentTypeError('File must have a csv or tab extension')
    return param

然后使用这个函数，比如：

parser = argparse.ArgumentParser()
parser.add_argument('filename', type=valid_file)

Answer 3

当然了，你只需要指定一个合适的函数作为 type。

import argparse
import os.path

parser = argparse.ArgumentParser()

def file_choices(choices,fname):
    ext = os.path.splitext(fname)[1][1:]
    if ext not in choices:
       parser.error("file doesn't end with one of {}".format(choices))
    return fname

parser.add_argument('fn',type=lambda s:file_choices(("csv","tab"),s))

parser.parse_args()

示例：

temp $ python test.py test.csv
temp $ python test.py test.foo
usage: test.py [-h] fn
test.py: error: file doesn't end with one of ('csv', 'tab')

---

这里有一种可能更简洁/通用的方法：

import argparse
import os.path

def CheckExt(choices):
    class Act(argparse.Action):
        def __call__(self,parser,namespace,fname,option_string=None):
            ext = os.path.splitext(fname)[1][1:]
            if ext not in choices:
                option_string = '({})'.format(option_string) if option_string else ''
                parser.error("file doesn't end with one of {}{}".format(choices,option_string))
            else:
                setattr(namespace,self.dest,fname)

    return Act

parser = argparse.ArgumentParser()
parser.add_argument('fn',action=CheckExt({'csv','txt'}))

print parser.parse_args()

不过这样做的缺点是，代码在某些方面会变得有点复杂。不过好处是，当你真正去格式化你的参数时，界面会变得更整洁。