python argparse.FileType('w') check extension

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]$'))

Solution

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.