python argparse.FileType ('w' ) 检查扩展名

Question

argparse 包在处理命令行参数时做得很好。但是我想知道是否有任何方法可以让 argparse 检查文件扩展名(例如“.txt”)。这个想法是派生一个与 argparse.FileType 相关的类。我会对任何建议感兴趣。

请记住，我的程序中有超过 50 个子命令，它们都有自己的 CLI。因此，与在我的所有命令中添加一些丑陋的测试相比，我更感兴趣的是派生一个可以在每个命令中导入的类。

非常感谢。

# 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

您可以将 argparse.FileType() class 子类化，并覆盖 __call__ 方法来进行文件名验证:

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)

如果您真的愿意，您也可以支持正则表达式，但使用 str.endswith() 是一种更常见和更简单的测试。

这需要单个字符串或指定有效扩展名的字符串元组:

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

您需要在 __call__ 方法中处理它，因为 FileType() 实例本质上与任何其他 type= 参数一样对待；作为 callable，您可以通过引发 ArgumentTypeError 异常来指示特定参数不合适。