Ask user to choose from a list of choices via command line

ask_user_for_selection.py

"""
Allows choice of options via command line. 

Author: awan1
Date: 20151029
"""

import numpy as np

def ask_user_for_selection(choices, choice_text='options'):
    """
    Given a list of choices, ask the user to pick a subset via command-line
    input.

    Parameters:
    choices: a list of choices to select from
    choice_text: text explaining what the choices are.

    Returns:
    The items in choices that the user specified

    """
    # Convert to array to allow selection via list of indices
    choices_array = np.array(choices)

    print('Please choose from the following {} by index. '
          'Multiple choices are allowed.'.format(choice_text))
    for i, choice in enumerate(choices):
        print("{}) {}".format(i, choice))
    all_option = i+1
    print("{}) All of the above".format(all_option))

    # Until we get valid input, try to get a set of numbers from user.
    while True:
        ans = raw_input('Choices: ')
        try:
            # Convert to int and into array to ease operations later
            # (checking that all elements are valid, masking the returns)
            ans_array = np.array([int(x) for x in ans.split()])
            if all(ans_array <= all_option) and all(ans_array >= 0):
                break
        except ValueError:
            pass
        # If we reach here, we didn't get good input.
        print("Invalid input. Try again.")

    # If the all_option is in the answers, we know to return everything.
    if all_option in ans_array:
        return choices
    else:
        return list(choices_array[ans_array])