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实现这一点？

- bph

Answer 1

当然可以 - 您只需要将适当的函数指定为“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()

这里的缺点是代码在某些方面变得更加复杂——好处是当您实际格式化参数时，界面会变得更加整洁。

- mgilson

---

嗯...... notcsv 作为输入名称是不正确的吧？这还取决于 parser 是否为全局变量，我认为引发异常才是正确的方式... - Jon Clements

我会说主要的区别在于这种方法需要将parser作为函数范围内的全局变量（因为它使用了parser.error），我避免了这种情况，而是从模块中引发了异常... - Jon Clements

@Hiett -- 它允许我使用相同的函数来检查多个命令行参数（例如，如果您有一个需要检查“.txt”或“.csv”的参数和另一个需要检查“.xls”或“.doc”的参数-- 否则，它们之间没有太大区别。 - mgilson

@mgilson 也是这样，我猜...;) 谁想再写一个答案来做出最终决定？:) - Jon Clements

@JonClements -- 我写了它。为什么不呢？ - mgilson

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)

- Jon Clements

---

你会在 parser.parse_args() 之前调用 valid_file() 吗？ - bph

Answer 3

不可以。你可以将一个容器对象提供给choices参数，或者任何支持“in”运算符的东西。你可以在pydocs中阅读更多信息。

不过你可以随时自己检查并向用户提供反馈。

- tayfun