Alias for argparse choice

Viewed 786

Is it possible to introduce an alias for a choice in an argparse argument?

Specifically, I want to let the user specify a short argument such as main.py --task a and automatically expand that to {'task': "a_long_filename"} after running parse_args().

I was hoping something along the following lines would exist but couldn't find anything in the argparse docs:

parser.add_argument("--task",
                    choices={"a": "a_long_filename", "b": "b_long_filename"},
                    default=""a_long_filename"", help="Task to be solved."
                   )
2 Answers
import argparse


TASK_CHOICES = {"a": "a_long_filename", "b": "b_long_filename"}


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--task",
                    choices=TASK_CHOICES.keys(),
                    default="a", 
                    help="Task to be solved.",
                   )
    return parser, parser.parse_args()


parser, args = get_args()
args.task = TASK_CHOICES.get(args.task)
print(args)

Outputs:

$ python3 arg_task.py --task a
Namespace(task='a_long_filename')
$ python3 arg_task.py --task b
Namespace(task='b_long_filename')

You can implement a subclass of argparse.Action which will do the conversion:

import argparse

parser = argparse.ArgumentParser()

choice_mapping = {"a": "a_long_filename", "b": "b_long_filename"}


class MyAction(argparse.Action):
    def __call__(self, parser, namespace, values, option_string):
        setattr(namespace, self.dest, choice_mapping[values[0]])


parser.add_argument(
    "--task",
    choices=("a", "b"),

    # your action won't be called on default value,
    # so changed from "a" to "a_long_filename"
    default="a_long_filename",  

    help="Task to be solved",
    action=MyAction
)

args = parser.parse_args(['--task', 'b'])

print(args)

# Namespace(task='b_long_filename')
Related