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()该类的子类，并重写该__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异常来指示特定参数不合适。

Answer 2

我的解决方案是创建一个进行扩展检查的闭包：

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

笔记

• ext_check 基本上是一个包装 argparse.FileType
• 需要预期的扩展名才能进行检查并需要打开
• 为简单起见，预期的扩展名使用小写，验证之前文件名将转换为小写
• openner在这种情况下是argparse.FileType('w')可调用的（很可能是一个函数，但是我不在乎，只要它是可调用的即可）。

• ext_check返回一个可调用的函数，该函数称为extension。我以这种方式命名，这样错误就会如下所示（请注意扩展名“波纹管”，它是函数的名称）：

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

• 在extension函数中，我们检查文件扩展名，如果通过，则将文件名传递给openner。

我喜欢这个解决方案

• 简洁
• 几乎不需要任何argparse.FileType工作原理的知识，因为它只是围绕它的包装

我对此不满意

• 呼叫者必须了解闭包以了解其工作原理
• 我无法控制该错误消息。这就是为什么我必须命名我的内部函数extension以获得如上所述的有意义的错误消息的原因。

其他可能的解决方案

• 创建自定义操作，请参阅argparse的文档
• argparse.FileTypeMartijn Pieters完成的子类

每个解决方案都有其优点和缺点