使用argparse指定文件扩展名

Question

我想使用argparse指定几个文件扩展名。

我尝试了下面的代码，但它不起作用。如何使用argparse指定多个文件扩展名？

parser.add_argument('file', action = 'store', type = argparse.FileType('r'), choices=["*.aaa", "*.bbb"])

编辑：我已经找到了自己的解决方案，使用字符串类型而不是FileType：

def input_ok(string):
    if not os.path.exists(string):
        raise argparse.ArgumentTypeError("Filename %r doesn\'t exists in this directory." % string)

    if string[-4:] != ".aaa" and string[-4:] != ".bbb":
        raise argparse.ArgumentTypeError("%r is not a .aaa or a .bbb file." % string)
return string

...

parser.add_argument('input_path', action = 'store',
    type = input_ok, #argparse.FileType('r'), #choices=["*.stl", "*.csv"])

- roipoussiere

Answer 1

问题的症结在于choices如何工作。Argparse首先构建传递参数的列表并进行类型转换，然后使用in运算符检查是否包含在选项中。这意味着它不会执行任何模式匹配（以匹配'*.aaa'），而是检查字符串相等性。相反，我们可以创建自己的容器传递给choices。

如果不使用argparse.Filetype，则如下所示。Argparse还需要该容器具有__iter__才能为帮助消息创建元变量元组。

class Choices():
    def __init__(self, *choices):
        self.choices = choices

    def __contains__(self, choice):
        # True if choice ends with one of self.choices
        return any(choice.endswith(c) for c in self.choices)

    def __iter__(self):
        return iter(self.choices)

parser.add_argument('file', action='store', choices=Choices('.aaa', '.bbb'))

您可以通过更改 __contains__ 来扩展此想法以执行几乎任何操作。例如，如果还传递了 type = argparse.FileType（'r'），则argparse将在检查包含之前转换为文件对象。

def __contains__(self, choice):
    # choice is now a file object
    return any(choice.name.endswith(c) for c in self.choices)

---

作为一个旁白，这就是我为什么讨厌argparse。它过于复杂，并试图做比它应该做的更多。我认为验证不应该在参数解析器中完成。使用docopt并自行验证。

- kalhartt

---

argparse 为参数验证提供了一些工具，但不要求您使用它们。 FileType 是一个方便的功能，旨在用于常见的脚本应用程序。如果不适合您的应用程序，则无需使用它。choices 同样如此。 - hpaulj

Answer 2

argparse 接受字符串并在此之后进行验证没有任何问题。有时候，您可能想要检查一个文件名是否正确，但在稍后不打开它（例如，使用 with open(filename) as f:）。这很可能是最简单的方法。

与 kalhartt 的 Choices 类相比，另一种选择是使用 os.path 或 glob 获取可接受的文件列表。

p.add_argument('file',choices=glob.glob('*.txt'))
In [91]: p.parse_args('test.txt'.split())
Out[91]: Namespace(file='test.txt')

这种方法的问题在于帮助和错误信息可能会过长，列出所有可允许的文件名。

此choices无法与FileType一起使用。这是因为它在文件打开后测试选择。

p.add_argument('file',choices=[open('test.txt')],type=argparse.FileType('r'))
p.parse_args('test.txt'.split())
# usage: python [-h] {<open file 'test.txt', mode 'r' at 0xa102f98>}
# error: argument file: invalid choice: <open file 'test.txt', mode 'r' at 0xa102f40>
# (choose from <open file 'test.txt', mode 'r' at 0xa102f98>)

即使文件名相同，两个打开的文件的ID也不相同。正如kalhartt的示例所示，choices对象将需要具有自定义__contains__函数（测试文件名的函数，例如f.name.endswith（'txt'））。

但是，如果您真的喜欢FileType打开文件的事实，我可以想象它的子类化，因此它会检查扩展名。

class FileTypeWithExtension(argparse.FileType):
    def __init__(self, mode='r', bufsize=-1, extension=None):
        self._extension = extension
        super(FileTypeWithExtension, self).__init__()
    def __call__(self, string):
        if string != '-' and self._extension:
            if not string.endswith(self._extension):
               # just testing against one extension for now
               raise argparse.ArgumentTypeError('wrong extension')
        return super(FileTypeWithExtension, self).__call__(string)

p.add_argument('file',type=FileTypeWithExtension('r',extension='txt'))
p.parse_args('test.tst'.split())
#usage: ipython [-h] file
#ipython: error: argument file: wrong extension

p.parse_args('test.txt'.split())
# Namespace(file=<open file 'test.txt', mode 'r' at 0xa13ce90>)

- hpaulj