用argparse指定文件擴展名

Question

我想用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"]) 

Answer 1

問題的關鍵是如何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並自行驗證事物。

Answer 2

​​接受一個字符串沒有問題，你在做自己的驗證之後。有時候你想檢查一個文件名是否正確，但直到以後纔打開它（例如使用with open(filename) as f:）。這可能是最簡單的方法。

到kalhartt'sChoices類另一種方法是使用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') 

這樣做的一個問題是，幫助和錯誤消息可能過長，列出所有允許的文件名。

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