python argparse.FileType('w') check extension

Question

The argparse package does a great job when dealing with command line arguments. I'm wondering however if there is any way to ask argparse to check for file extension (e.g ".txt"). The idea would be to derived one the class related to argparse.FileType. I would be interested in any suggestion.

Keep in mind that I have more than 50 subcommands in my program all having there own CLI. Thus, I would be interest in deriving a class that could be imported in each of them more than adding some uggly tests in all my commands.

Thanks a lot.

# As an example one would be interested in turning this...
parser_grp.add_argument('-o', '--outputfile',
                        help="Output file.",
                        default=sys.stdout,
                        metavar="TXT",
                        type=argparse.FileType('w'))


# Into that...
from somewhere import FileTypeWithExtensionCheck 
parser_grp.add_argument('-o', '--outputfile',
                        help="Output file.",
                        default=sys.stdout,
                        metavar="TXT",
                        type=FileTypeWithExtensionCheck('w', '.[Tt][Xx][Tt]$'))

Answer 1

You could subclass the argparse.FileType() class, and override the __call__ method to do filename validation:

class FileTypeWithExtensionCheck(argparse.FileType):
    def __init__(self, mode='r', valid_extensions=None, **kwargs):
        super().__init__(mode, **kwargs)
        self.valid_extensions = valid_extensions

    def __call__(self, string):
        if self.valid_extensions:
            if not string.endswith(self.valid_extensions):
                raise argparse.ArgumentTypeError(
                    'Not a valid filename extension')
        return super().__call__(string)

You could also support a regex if you really want to, but using str.endswith() is a more common and simpler test.

This takes either a single string, or a tuple of strings specifying valid extensions:

parser_grp.add_argument(
    '-o', '--outputfile', help="Output file.",
    default=sys.stdout, metavar="TXT",
    type=argparse.FileTypeWithExtensionCheck('w', valid_extensions=('.txt', '.TXT', '.text'))
)

You need to handle this in the __call__ method because the FileType() instance is essentially treated like any other type= argument; as a callable, and you can indicate that the specific argument isn't suitable by raising the ArgumentTypeError exception.

Answer 2

My solution is to create an closure that does the extension checking:

import argparse

def ext_check(expected_extension, openner):
    def extension(filename):
        if not filename.lower().endswith(expected_extension):
            raise ValueError()
        return openner(filename)
    return extension

parser = argparse.ArgumentParser()
parser.add_argument('outfile', type=ext_check('.txt', argparse.FileType('w')))

# test out
args = parser.parse_args()
args.outfile.write('Hello, world\n')

Notes

• ext_check basically is a wrapper for argparse.FileType
• It takes an expected extension to check and an openner
• For simplicity, the expected extension is in lower case, the filename will be converted to lower case prior to validation
• openner in this case is an argparse.FileType('w') callable (most likely a function, but I don't care, as long as it is a callable).

• ext_check returns a callable, which is a function called extension. I name it this way, so that the error will come out as followed (note the word extension bellow, which is the name of the function):

  error: argument outfile: invalid extension value: 'foo.txt2'
  

• Within the extension function, we check the file extension, if passed, we pass the file name to the openner.

What I like about this solution

• Concise
• Require almost no knowledge of how argparse.FileType works since it just act as a wrapper around it

What I don't like about it

• Caller has to know about closure to understand how it works
• I have no control over the error message. That is why I have to name my inner function extension to get a somewhat meaningful error message as seen above.

Other possible solutions

• Create a custom action, see the documentation for argparse
• Subclass argparse.FileType as Martijn Pieters has done

Each of these solutions has its own strong points and weaknesses