WIP Chromite supporting "LEGACY" mode, as requested by Anush.

chromite/lib/text_menu.py

+#!/usr/bin/python
+# 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 __ceil_div(numerator, denominator):
+  """Do integer division, rounding up.
+
+  >>> __ceil_div(-1, 2)
+  0
+  >>> __ceil_div(0, 2)
+  0
+  >>> __ceil_div(1, 2)
+  1
+  >>> __ceil_div(2, 2)
+  1
+  >>> __ceil_div(3, 2)
+  2
+
+  Args:
+    numerator: The number to divide.
+    denominator: The number to divide by.
+
+  Returns:
+    (numberator) / denominator, rounded up.
+  """
+  return (numerator + denominator - 1) // denominator
+
+
+def __build_menu_str(items, title, prompt, menu_width, spacing):
+  """Build the menu string for TextMenu.
+
+  See TextMenu for a description.  This function mostly exists to simplify
+  testing.
+
+  >>> __build_menu_str(['A'], "Choose:", "Choice: ", 76, 2)
+  'Choose:\\n\\n1. A\\n\\nChoice: '
+
+  >>> __build_menu_str(['B', 'A'], "Choose:", "Choice: ", 76, 2)
+  'Choose:\\n\\n1. B  2. A\\n\\nChoice: '
+
+  >>> __build_menu_str(['A', 'B', 'C', 'D', 'E'], "Choose:", "Choice: ", 10, 2)
+  'Choose:\\n\\n1. A  4. D\\n2. B  5. E\\n3. C\\n\\nChoice: '
+
+  >>> __build_menu_str(['A', 'B', 'C', 'D', 'E'], "Choose:", "Choice: ", 11, 2)
+  'Choose:\\n\\n1. A  4. D\\n2. B  5. E\\n3. C\\n\\nChoice: '
+
+  >>> __build_menu_str(['A', 'B', 'C'], "Choose:", "Choice: ", 9, 2)
+  'Choose:\\n\\n1. A\\n2. B\\n3. C\\n\\nChoice: '
+
+  >>> __build_menu_str(['A'*10, 'B'*10], "Choose:", "Choice: ", 0, 2)
+  'Choose:\\n\\n1. AAAAAAAAAA\\n2. BBBBBBBBBB\\n\\nChoice: '
+
+  >>> __build_menu_str([], "Choose:", "Choice: ", 76, 2)
+  Traceback (most recent call last):
+  ...
+  ValueError: Can't build a menu of empty choices.
+  """
+  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 = __ceil_div(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):
+    col, row = divmod(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]
+
+  # 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):
+  """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, it is expected that he/she
+  will press Ctrl-C to exit.
+
+  The menu will look something like this:
+    1. __init__.py   3. chromite      5. lib           7. tests
+    2. bin           4. chroot_specs  6. specs
+
+    Choice:
+
+  Args:
+    items: The strings to show in the menu.  These should be sorted in whatever
+        order you want to show to the user.
+    prompt: The prompt to show to the user.
+    menu_width: The maximum width to use for the menu; 0 forces things to single
+        column.
+    spacing: The spacing between items.
+
+  Returns:
+    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.
+  """
+  # Call the helper to build the actual menu string.
+  menu_str = __build_menu_str(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.
+    sys.stderr.write(menu_str)
+
+    # Don't use a prompt with raw_input(), since that would go to stdout.
+    result = raw_input()
+
+    # Parse into a number and do error checking.  If all good, return.
+    try:
+      result_int = int(result)
+      if 1 <= result_int <= len(items):
+        # Convert from 1-based to 0-based index!
+        return result_int - 1
+      else:
+        print >>sys.stderr, "\nERROR: %d out of range\n" % result_int
+    except ValueError:
+      print >>sys.stderr, "\nERROR: '%s' is not a valid choice\n" % result
+
+
+def __test():
+  """Run any built-in tests."""
+  import doctest
+  doctest.testmod(verbose=True)
+
+
+# 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']:
+    __test(*sys.argv[2:])
+  else:
+    TextMenu(sys.argv[1:])