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

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.

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)

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.

来源：https://stackoverflow.com/questions/15203829/python-argparse-file-extension-checking

标签

python

argparse