Best way to handle if-else with argparse?

Say we have a lot of options for argument parsing

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Download files from Canvas.')

    parser.add_argument(
        '-user_id',
        type=str,
        default="",
    )

    parser.add_argument(
        '-course_id',
        type=str,
        default="",
    )

    parser.add_argument(
        '-student_id',
        type=str,
        default="",
    )
# ...
    parser.add_argument(
        '-book_id',
        type=str,
        default="",
    )

And I want to run functions determined by the inputs

if user_id != "" and course_id == "" and student_id != "" and book_id == "":
    f()

elif user_id != "" and course_id == "" and student_id != "" and book_id != "":
    g()

# ...

elif user_id == "" and course_id == "" and student_id == "" and book_id == "":
    h()

and so on. Well, these are a lot of if-else statements. I want to eliminate having to do all of that. Does anyone have sufficient ways to reduce the amount of comparisons among variables before a function is run? Or does it depend on context? Thank you!!

Answer

Build a dictionary that maps a combination of options to a handler function.

When you write this:

if user_id != "" and course_id == "" and student_id != "" and book_id == "":

You are effectively treating those four options as a tuple of boolean values. You are asking, “are the arguments equal to (False, True, False, True)?”

We can take advantage of that. First, we need code to translate the options into an actual tuple. Something like:

args = parser.parse_args()
sig = tuple(getattr(args, x) == '' for x in ('user_id', 'course_id', 'student_id', 'book_id'))

If your program receives no command line arguments, this sets sig equal to (True, True, True, True). If you pass --student-id 1, this becomes (True, True, False, True).

We can create a dictionary that maps these patterns to a handler function:

argmap = {
    (False, True, False, True): f,
    (False, True, False, False): g,
    (True, True, True, True): h,
}

With these two mechanisms in hand, we can look up the handler for a given set of arguments like this:

handler = argmap.get(sig)
if handler is None:
    print('no handler for', args)
else:
    handler()

Putting it all together, we get:

import argparse


def f():
    print('this is function f')


def g():
    print('this is function g')


def h():
    print('this is function h')


argmap = {
    (False, True, False, True): f,
    (False, True, False, False): g,
    (True, True, True, True): h,
}


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Download files from Canvas.')

    parser.add_argument(
        '--user-id',
        default='',
    )

    parser.add_argument(
        '--course-id',
        default='',
    )

    parser.add_argument(
        '--student-id',
        default='',
    )

    parser.add_argument(
        '--book-id',
        default='',
    )

    args = parser.parse_args()
    sig = tuple(getattr(args, x) == '' for x in ('user_id', 'course_id', 'student_id', 'book_id'))
    handler = argmap.get(sig)
    if handler is None:
        print('no handler for', args)
    else:
        handler()

This produces as output:

$ python argtest.py
this is function h
$ python argtest.py --user-id 1 --student-id 1
this is function f
$ python argtest.py --user-id 1 --student-id 1 --book-id 1
this is function g
$ python argtest.py --user-id 1
no handler for Namespace(user_id='1', course_id='', student_id='', book_id='')

Leave a Reply

Your email address will not be published. Required fields are marked *