Specify file extension with argparse

Question

I would like to specify several file extensions with argparse.

I have tried the code below but it doesn't work. How do I specify multiple file extensions with argparse?

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

EDIT : I have found my own solution using string type instead of 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

The crux of the problem is how choices work. Argparse first builds a list of the passed arguments and does type conversion, then it checks for inclusion in choices with the in operator. This means it won't do any pattern matching (to match '*.aaa') but checks for string equality. Instead we can make our own container to pass to choices.

Without using argparse.Filetype it looks like this. Argparse also needs the container to have __iter__ to make the Metavar tuple for the help message.

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

You can extend this idea to do pretty much anything by changing __contains__ to suit your needs. For example, if you also passed type=argparse.FileType('r') then argparse will convert to a file object before checking inclusion.

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

---

As an aside, this is why I hate argparse. It's overly complex and tries to do way more than it should. I don't think validation should be done in the argument parser. Use docopt and validate things yourself.

Answer 2

There's nothing wrong with argparse accepting a string, and you doing your own validation after. Sometimes you want to check if a filename is correct, but not open it until later (e.g. using with open(filename) as f:). This is likely to be the simplest method.

An alternative to kalhartt's Choices class would be to use os.path or glob to get a list of allowable files.

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

A problem with this is that the help and error messages can be excessively long, listing all of the allowable filenames.

This choices does not work along with FileType. That's because it tests against choices after the file has been opened

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

Even though the filenames are the same, the ids of the two opened files are not the same. Askalhartt's example shows, the choices object will have to have a custom __contains__ function (one that tests the file name, e.g. f.name.endswith('txt')).

But if you really like the fact that FileType opens the file, I can imagine subclassing it, so it checks for extensions.

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