python argparse file extension checking

Question

can argparse be used to validate filename extensions for a filename cmd line parameter?

e.g. if i have a python script i run from the cmd line:

Answer 1

Sure -- you just need to specify an appropriate function as the type.

import argparse
import os.path

parser = argparse.ArgumentParser()

def file_choices(choices,fname):
    ext = os.path.splitext(fname)[1][1:]
    if ext not in choices:
       parser.error("file doesn't end with one of {}".format(choices))
    return fname

parser.add_argument('fn',type=lambda s:file_choices(("csv","tab"),s))

parser.parse_args()

demo:

temp $ python test.py test.csv
temp $ python test.py test.foo
usage: test.py [-h] fn
test.py: error: file doesn't end with one of ('csv', 'tab')

---

Here's a possibly more clean/general way to do it:

import argparse
import os.path

def CheckExt(choices):
    class Act(argparse.Action):
        def __call__(self,parser,namespace,fname,option_string=None):
            ext = os.path.splitext(fname)[1][1:]
            if ext not in choices:
                option_string = '({})'.format(option_string) if option_string else ''
                parser.error("file doesn't end with one of {}{}".format(choices,option_string))
            else:
                setattr(namespace,self.dest,fname)

    return Act

parser = argparse.ArgumentParser()
parser.add_argument('fn',action=CheckExt({'csv','txt'}))

print parser.parse_args()

The downside here is that the code is getting a bit more complicated in some ways -- The upshot is that the interface gets a good bit cleaner when you actually go to format your arguments.

Answer 2

Define a custom function which takes the name as a string - split the extension off for comparison and just return the string if it's okay, otherwise raise the exception that argparse expects:

def valid_file(param):
    base, ext = os.path.splitext(param)
    if ext.lower() not in ('.csv', '.tab'):
        raise argparse.ArgumentTypeError('File must have a csv or tab extension')
    return param

And then use that function, such as:

parser = argparse.ArgumentParser()
parser.add_argument('filename', type=valid_file)

Answer 3

No. You can provide a container object to choices argument, or anything that supports the "in" operator. You can read more at pydocs

You can always check it yourself and provide feedback to the user though.