Python Argparse File Extension Checking

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 fil

Solution 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 notin 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] fntest.py: error: filedoesn'tendwithoneof ('csv', 'tab')

---

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

import argparse
import os.path

defCheckExt(choices):
    classAct(argparse.Action):
        def__call__(self,parser,namespace,fname,option_string=None):
            ext = os.path.splitext(fname)[1][1:]
            if ext notin 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.

Solution 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() notin ('.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)

Solution 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.