Created
November 8, 2018 21:00
-
-
Save jesteria/638c719875d38023135966d0aad78937 to your computer and use it in GitHub Desktop.
argparse action to store value defined in dict
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import argparse | |
| class StoreMapAction(argparse.Action): | |
| Unspec = object() | |
| def __init__(self, | |
| option_strings, | |
| dest, | |
| map, | |
| nargs=None, | |
| default=Unspec, | |
| type=None, | |
| choices=Unspec, | |
| required=False, | |
| help=None, | |
| metavar=None): | |
| self.choicemap = map if hasattr(map, 'values') else dict(map) | |
| if choices is self.Unspec: | |
| choices = self.choicemap | |
| if choices and default is self.Unspec: | |
| try: | |
| choice_values = choices.values() | |
| except AttributeError: | |
| choice_values = choices | |
| default = next(iter(choice_values)) | |
| super().__init__( | |
| option_strings=option_strings, | |
| dest=dest, | |
| nargs=nargs, | |
| const=None, | |
| default=default, | |
| type=type, | |
| choices=choices, | |
| required=required, | |
| help=help, | |
| metavar=metavar, | |
| ) | |
| def __call__(self, parser, namespace, values, option_string=None): | |
| map_value = self.choicemap.get(values, values) | |
| setattr(namespace, self.dest, map_value) | |
| ### | |
| types = { | |
| 'aaa': 'AAA', | |
| 'bbb': 'BBB', | |
| } | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument( | |
| '--type', | |
| action=StoreMapAction, | |
| map=types, | |
| ) | |
| ### | |
| for (argv, namespace) in ( | |
| ((), argparse.Namespace(type='AAA')), | |
| (('--type', 'aaa'), argparse.Namespace(type='AAA')), | |
| (('--type', 'bbb'), argparse.Namespace(type='BBB')), | |
| ): | |
| print(argv, namespace) | |
| args = parser.parse_args(argv) | |
| assert args == namespace | |
| try: | |
| parser.parse_args(('--type', 'BBB')) | |
| except BaseException: | |
| print("it worked") | |
| else: | |
| raise AssertionError |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment