Укажите расширение файла с помощью argparse

Question

Я хотел бы указать несколько расширений файлов с помощью argparse.

Я пробовал код ниже, но он не работает. Как указать несколько расширений файлов с помощью argparse?

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

EDIT: Я нашел свое собственное решение, используя тип строки вместо 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"]) 

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

Нет ничего плохого в том, что argparse принимает строку, и вы выполняете свою собственную проверку после. Иногда вы хотите проверить правильность имени файла, но не открывать его до конца (например, с помощью 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') 

Проблема с этим состоит в том, что помощь и сообщения об ошибках могут быть чрезмерно длинным списком всех допустимых имен файлов.

Этот 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>) 

Даже если имена файлов совпадают, идентификаторы двух открытых файлов не совпадают. Как показано в примере 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>)