blob: aa6d0446f8cf649e2c867da5dbc90e2dfe7e03bf [file] [log] [blame]
# Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""A module containing a function for presenting a text menu to the user."""
import sys
def _CeilDiv(numerator, denominator):
"""Do integer division, rounding up.
>>> _CeilDiv(-1, 2)
>>> _CeilDiv(0, 2)
>>> _CeilDiv(1, 2)
>>> _CeilDiv(2, 2)
>>> _CeilDiv(3, 2)
numerator: The number to divide.
denominator: The number to divide by.
(numberator) / denominator, rounded up.
return (numerator + denominator - 1) // denominator
def _BuildMenuStr(items, title, prompt, menu_width, spacing, add_quit):
"""Build the menu string for TextMenu.
See TextMenu for a description. This function mostly exists to simplify
>>> _BuildMenuStr(['A'], "Choose", "Choice", 76, 2, True)
'Choose:\\n\\n1. A\\n\\nChoice (q to quit): '
>>> _BuildMenuStr(['B', 'A'], "Choose", "Choice", 76, 2, False)
'Choose:\\n\\n1. B 2. A\\n\\nChoice: '
>>> _BuildMenuStr(['A', 'B', 'C', 'D', 'E'], "Choose", "Choice", 10, 2, False)
'Choose:\\n\\n1. A 4. D\\n2. B 5. E\\n3. C\\n\\nChoice: '
>>> _BuildMenuStr(['A', 'B', 'C', 'D', 'E'], "Choose", "Choice", 11, 2, False)
'Choose:\\n\\n1. A 4. D\\n2. B 5. E\\n3. C\\n\\nChoice: '
>>> _BuildMenuStr(['A', 'B', 'C'], "Choose", "Choice", 9, 2, False)
'Choose:\\n\\n1. A\\n2. B\\n3. C\\n\\nChoice: '
>>> _BuildMenuStr(['A'*10, 'B'*10], "Choose", "Choice", 0, 2, False)
'Choose:\\n\\n1. AAAAAAAAAA\\n2. BBBBBBBBBB\\n\\nChoice: '
>>> _BuildMenuStr([], "Choose", "Choice", 76, 2, False)
Traceback (most recent call last):
ValueError: Can't build a menu of empty choices.
items: See TextMenu().
title: See TextMenu().
prompt: See TextMenu().
menu_width: See TextMenu().
spacing: See TextMenu().
add_quit: See TextMenu().
See TextMenu().
ValueError: If no items.
if not items:
raise ValueError("Can't build a menu of empty choices.")
# Figure out some basic stats about the items.
num_items = len(items)
longest_num = len(str(num_items))
longest_item = longest_num + len(". ") + max(len(item) for item in items)
# Figure out number of rows / cols.
num_cols = max(1, (menu_width + spacing) // (longest_item + spacing))
num_rows = _CeilDiv(num_items, num_cols)
# Construct "2D array" of lines. Remember that we go down first, then
# right. This seems to mimic "ls" behavior. Note that, unlike "ls", we
# currently make all columns have the same width. Shrinking small columns
# would be a nice optimization, but complicates the algorithm a bit.
lines = [[] for _ in xrange(num_rows)]
for item_num, item in enumerate(items):
row = item_num % num_rows
item_str = "%*d. %s" % (longest_num, item_num + 1, item)
lines[row].append("%-*s" % (longest_item, item_str))
# Change lines from 2D array into 1D array (1 entry per row) by joining
# columns with spaces.
spaces = " " * spacing
lines = [spaces.join(line) for line in lines]
# Add '(q to quit)' string to prompt if requested...
if add_quit:
prompt = "%s (q to quit)" % prompt
# Make the final menu string by adding some return and the prompt.
menu_str = "%s:\n\n%s\n\n%s: " % (title, "\n".join(lines), prompt)
return menu_str
def TextMenu(items, title="Choose one", prompt="Choice",
menu_width=76, spacing=4, add_quit=True):
"""Display text-based menu to the user and get back a response.
The menu will be printed to sys.stderr and input will be read from sys.stdin.
If the user doesn't want to choose something, he/she can use the 'q' to quit
or press Ctrl-C (which will be caught and treated the same).
The menu will look something like this:
1. 3. chromite 5. lib 7. tests
2. bin 4. chroot_specs 6. specs
Choice (q to quit):
items: The strings to show in the menu. These should be sorted in whatever
order you want to show to the user.
title: The title of the menu.
prompt: The prompt to show to the user.
menu_width: The maximum width to use for the menu; 0 forces things to single
spacing: The spacing between items.
add_quit: Let the user type 'q' to quit the menu (we'll return None).
The index of the item chosen by the user. Note that this is a 0-based
index, even though the user is presented with the menu in 1-based format.
Will be None if the user hits Ctrl-C, or chooses q to quit.
ValueError: If no items.
# Call the helper to build the actual menu string.
menu_str = _BuildMenuStr(items, title, prompt, menu_width, spacing,
# Loop until we get a valid input from the user (or they hit Ctrl-C, which
# will throw and exception).
while True:
# Write the menu to stderr, which makes it possible to use this with
# commands where you want the output redirected.
# Don't use a prompt with raw_input(), since that would go to stdout.
result = raw_input()
except KeyboardInterrupt:
# Consider this a quit.
return None
# Check for quit request
if add_quit and result.lower() in ("q", "quit"):
return None
# Parse into a number and do error checking. If all good, return.
result_int = int(result)
if 1 <= result_int <= len(items):
# Convert from 1-based to 0-based index!
return result_int - 1
print >>sys.stderr, "\nERROR: %d out of range.\n\n" % result_int
except ValueError:
print >>sys.stderr, "\nERROR: '%s' is not a valid choice.\n\n" % result
def _Test():
"""Run any built-in tests."""
import doctest
# For testing purposes, you can run this on the command line...
if __name__ == "__main__":
# If first argument is --test, run testing code.
# ...otherwise, pass all arguments as the menu to show.
if sys.argv[1:2] == ["--test"]:
if not sys.argv[1:]:
print "ERROR: Need params to display as menu items"
result = TextMenu(sys.argv[1:])
print "You chose: '%s'" % result