Можно ли вызвать ошибку, если add_argument () вызывается с недопустимым типом? - PullRequest
1 голос
/ 03 февраля 2020

Я бы хотел запретить некоторым типам объявляться как аргументы 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

1 Ответ

2 голосов
/ 03 февраля 2020

Вы можете подкласс 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
...