2014-01-05 54 views
3

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

我试过下面的代码,但它不起作用。如何使用argparse指定多个文件扩展名?

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

编辑:我发现使用字符串类型,而不是文件类型的我自己的解决方案:

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"]) 

回答

3

问题的关键是如何choices工作。 Argparse首先建立一个传递参数的列表并进行类型转换,然后它检查包含在in运算符的选项中。这意味着它不会进行任何模式匹配(匹配'*.aaa'),但会检查字符串是否相等。相反,我们可以让自己的容器通过选择。

没有使用argparse.Filetype它看起来像这样。 Argparse还需要容器使__iter__为帮助消息创建Metavar元组。

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并自行验证事物。

+1

'argparse'为您提供了一些参数验证工具,但并不要求您使用它们。 'FileType'是一个方便的功能,用于通用脚本应用程序。如果它不适合您的应用程序,则不必使用它。 “选择”也一样。 – hpaulj

1

​​接受一个字符串没有问题,你在做自己的验证之后。有时候你想检查一个文件名是否正确,但直到以后才打开它(例如使用with open(filename) as f:)。这可能是最简单的方法。

kalhartt'sChoices类另一种方法是使用os.pathglob获得的许可文件的列表。

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

这样做的一个问题是,帮助和错误消息可能过长,列出所有允许的文件名。

This choices does not work with 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>) 

后即使文件名是相同的,两者的IDS打开的文件是不一样的测试针对的选择。如kalhartt's示例所示,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>)