时间戳

首页 > 解决方案 > 如果使用不允许的类型调用 add_argument() 是否可能引发错误？

python - 如果使用不允许的类型调用 add_argument() 是否可能引发错误？

问题描述

我想禁止将某些类型声明为 argparse 参数。如果使用不允许的类型调用 add_argument() 是否可以禁止某些类型并引发错误？

伪代码例如：

parser = argparse.ArgumentParser(disallowed_types=[bool, str])

parser.add_argument('--seed', type=int, default=0) # Ok
parser.add_argument('--frac', type=float, default=0.1) # Ok
parser.add_argument('--flag', type=bool, default=False) # Raise an error
parser.add_argument('--name', type=str) # Raise an error

标签： pythonparsingargumentsargparse

解决方案

您可以子类化ArgumentParser.

import argparse

class MyParser(argparse.ArgumentParser):
    def __init__(self, disallowed_types, *args, **kwargs):
        self.disallowed_types = disallowed_types
        super().__init__(*args, **kwargs)

    def add_argument(self, *args, **kwargs):
        arg_type = kwargs.get('type')
        if arg_type in self.disallowed_types:
            raise TypeError('Using {} is disallowed'.format(arg_type))
        return super().add_argument(*args, **kwargs)

parser = MyParser(disallowed_types=(bool, str))

parser.add_argument('--seed', type=int, default=0)  # Ok
parser.add_argument('--frac', type=float, default=0.1)  # Ok
parser.add_argument('--flag', type=bool, default=False)  # Raise an error

Traceback (most recent call last):
File "main.py" in <module>
parser.add_argument('--flag', type=bool, default=False)  # Raise an error
File "main.py" in add_argument
raise TypeError('Using {} is disallowed'.format(arg_type))
TypeError: Using <class 'bool'> is disallowed

推荐阅读