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:

$ script.py file.csv
$ script.py file.tab
$ script.py file.txt

i would like argparse to accept the 1st two filename cmd line options but reject the 3rd

i know you can do something like this:

parser = argparse.ArgumentParser()
parser.add_argument("fn", choices=["csv","tab"])
args = parser.parse_args()

to specify two valid choices for a cmd line option

what i'd like is this:

parser.add_argument("fn", choices=["*.csv","*.tab"])

to specify two valid file extensions for the cmd line option. Unfortunately this doesn't work - is there a way to achieve this using argparse?

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)