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

Question

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

请记住，我的程序中有 50 多个子命令，它们都有自己的 CLI。因此，与在我的所有命令中添加一些难看的测试相比，我更感兴趣的是推导一个可以导入到每个命令中的 class。

非常感谢。

# 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 异常来指示特定参数不合适。

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 returns 一个可调用函数，它是一个名为 extension 的函数。我这样命名，这样报错就出来了（注意下面的extension这个词，是函数名）：

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

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

我喜欢这个解决方案的地方

• 简洁
• 几乎不需要了解 argparse.FileType 是如何工作的，因为它只是作为它的包装器

我不喜欢它的地方

• 调用者必须了解闭包才能理解其工作原理
• 我无法控制错误消息。这就是为什么我必须将我的内部函数命名为 extension 以获得如上所示的有点有意义的错误消息。

其他可能的解决方案

• 创建自定义操作，查看 argparse 的文档
• Martijn Pieters 所做的子类 argparse.FileType

这些解决方案中的每一个都有自己的优点和缺点