Split out the big chromite shell 'main.py' into lots of subfiles.

shell/utils.py

Index: shell/utils.py
diff --git a/shell/utils.py b/shell/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..8aef16c53de6055e979cbf8e81cf3765a07fb083
--- /dev/null
+++ b/shell/utils.py
@@ -0,0 +1,386 @@
+# Copyright (c) 2011 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.
+
+"""Utility functions shared between files in the chromite shell."""
+
+
+# Python imports
+import ConfigParser
+import cPickle
+import os
+import sys
+import tempfile
+
+
+# Local imports
+import chromite.lib.cros_build_lib as cros_lib
+from chromite.lib import text_menu
+
+
+# Find the Chromite root and Chromium OS root...  Note: in the chroot we may
+# choose to install Chromite somewhere (/usr/lib/chromite?), so we use the
+# environment variable to get the right place if it exists.
+CHROMITE_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), '..')
+SRCROOT_PATH = os.environ.get('CROS_WORKON_SRCROOT',
+                              os.path.realpath(os.path.join(CHROMITE_PATH,
+                                                            '..')))
+
+
+# Commands can take one of these two types of specs.  Note that if a command
+# takes a build spec, we will find the associated chroot spec.  This should be
+# a human-readable string.  It is printed and also is the name of the spec
+# directory.
+BUILD_SPEC_TYPE = 'build'
+CHROOT_SPEC_TYPE = 'chroot'
+
+
+# This is a special target that indicates that you want to do something just
+# to the host.  This means different things to different commands.
+# TODO(dianders): Good idea or bad idea?
+_HOST_TARGET = 'HOST'
+
+
+def GetBoardDir(build_config):
+  """Returns the board directory (inside the chroot) given the name.
+
+  Args:
+    build_config: A SafeConfigParser representing the config that we're
+        building.
+
+  Returns:
+    The absolute path to the board.
+  """
+  target_name = build_config.get('BUILD', 'target')
+
+  # Extra checks on these, since we sometimes might do a rm -f on the board
+  # directory and these could cause some really bad behavior.
+  assert target_name, "Didn't expect blank target name."
+  assert len(target_name.split()) == 1, 'Target name should have no whitespace.'
+
+  return os.path.join('/', 'build', target_name)
+
+
+def GetChrootAbsDir(chroot_config):
+  """Returns the absolute chroot directory the chroot config.
+
+  Args:
+    chroot_config: A SafeConfigParser representing the config for the chroot.
+
+  Returns:
+    The chroot directory, always absolute.
+  """
+  chroot_dir = chroot_config.get('CHROOT', 'path')
+  chroot_dir = os.path.join(SRCROOT_PATH, chroot_dir)
+
+  return chroot_dir
+
+
+def DoesChrootExist(chroot_config):
+  """Return True if the chroot folder exists.
+
+  Args:
+    chroot_config: A SafeConfigParser representing the config for the chroot.
+
+  Returns:
+    True if the chroot folder exists.
+  """
+  chroot_dir = GetChrootAbsDir(chroot_config)
+  return os.path.isdir(chroot_dir)
+
+
+def FindSpec(spec_name, spec_type=BUILD_SPEC_TYPE, can_show_ui=True):
+  """Find the spec with the given name.
+
+  This tries to be smart about helping the user to find the right spec.  See
+  the spec_name parameter for details.
+
+  Args:
+    spec_name: Can be any of the following:
+        1. A full path to a spec file (including the .spec suffix).  This is
+           checked first (i.e. if os.path.isfile(spec_name), we assume we've
+           got this case).
+        2. The full name of a spec file somewhere in the spec search path
+           (not including the .spec suffix).  This is checked second.  Putting
+           this check second means that if one spec name is a substring of
+           another, you can still specify the shorter spec name and know you
+           won't get a menu (the exact match prevents the menu).
+        3. A substring that will be used to pare-down a menu of spec files
+           found in the spec search path.  Can be the empty string to show a
+           menu of all spec files in the spec path.  NOTE: Only possible if
+           can_show_ui is True.
+    spec_type: The type of spec this is: 'chroot' or 'build'.
+    can_show_ui: If True, enables the spec name to be a substring since we can
+        then show a menu if the substring matches more than one thing.
+
+  Returns:
+    A path to the spec file.
+  """
+  spec_suffix = '.spec'
+
+  # Handle 'HOST' for spec name w/ no searching, so it counts as an exact match.
+  if spec_type == BUILD_SPEC_TYPE and spec_name == _HOST_TARGET.lower():
+    return _HOST_TARGET
+
+  # If we have an exact path name, that's it.  No searching.
+  if os.path.isfile(spec_name):
+    return spec_name
+
+  # Figure out what our search path should be.
+  # ...make these lists in anticipation of the need to support specs that live
+  # in private overlays.
+  # TODO(dianders): Should specs be part of the shell, or part of the main
+  # chromite?
+  search_path = [
+      os.path.join(CHROMITE_PATH, 'specs', spec_type),
+  ]
+
+  # Look for an exact match of a spec name.  An exact match will go through with
+  # no menu.
+  if spec_name:
+    for dir_path in search_path:
+      spec_path = os.path.join(dir_path, spec_name + spec_suffix)
+      if os.path.isfile(spec_path):
+        return spec_path
+
+  # cros_lib.Die right away if we can't show UI and didn't have an exact match.
+  if not can_show_ui:
+    cros_lib.Die("Couldn't find %s spec: %s" % (spec_type, spec_name))
+
+  # No full path and no exact match.  Move onto a menu.
+  # First step is to construct the options.  We'll store in a dict keyed by
+  # spec name.
+  options = {}
+  for dir_path in search_path:
+    for file_name in os.listdir(dir_path):
+      # Find any files that end with ".spec" in a case-insensitive manner.
+      if not file_name.lower().endswith(spec_suffix):
+        continue
+
+      this_spec_name, _ = os.path.splitext(file_name)
+
+      # Skip if this spec file doesn't contain our substring.  We are _always_
+      # case insensitive here to be helpful to the user.
+      if spec_name.lower() not in this_spec_name.lower():
+        continue
+
+      # Disallow the spec to appear twice in the search path.  This is the
+      # safest thing for now.  We could revisit it later if we ever found a
+      # good reason (and if we ever have more than one directory in the
+      # search path).
+      if this_spec_name in options:
+        cros_lib.Die('Spec %s was found in two places in the search path' %
+            this_spec_name)
+
+      # Combine to get a full path.
+      full_path = os.path.join(dir_path, file_name)
+
+      # Ignore directories or anything else that isn't a file.
+      if not os.path.isfile(full_path):
+        continue
+
+      # OK, it's good.  Store the path.
+      options[this_spec_name] = full_path
+
+  # Add 'HOST'.  All caps so it sorts first.
+  if (spec_type == BUILD_SPEC_TYPE and
+      spec_name.lower() in _HOST_TARGET.lower()):
+    options[_HOST_TARGET] = _HOST_TARGET
+
+  # If no match, die.
+  if not options:
+    cros_lib.Die("Couldn't find any matching %s specs for: %s" % (spec_type,
+                                                                  spec_name))
+
+  # Avoid showing the user a menu if the user's search string matched exactly
+  # one item.
+  if spec_name and len(options) == 1:
+    _, spec_path = options.popitem()
+    return spec_path
+
+  # If more than one, show a menu...
+  option_keys = sorted(options.iterkeys())
+  choice = text_menu.TextMenu(option_keys, 'Choose a %s spec' % spec_type)
+
+  if choice is None:
+    cros_lib.Die('OK, cancelling...')
+  else:
+    return options[option_keys[choice]]
+
+
+def ReadConfig(spec_path):
+  """Read the a build config or chroot config from a spec file.
+
+  This includes adding thue proper _default stuff in.
+
+  Args:
+    spec_path: The path to the build or chroot spec.
+
+  Returns:
+    config: A SafeConfigParser representing the config.
+  """
+  spec_name, _ = os.path.splitext(os.path.basename(spec_path))
+  spec_dir = os.path.dirname(spec_path)
+
+  config = ConfigParser.SafeConfigParser({'name': spec_name})
+  config.read(os.path.join(spec_dir, '_defaults'))
+  config.read(spec_path)
+
+  return config
+
+
+def GetBuildConfigFromArgs(argv):
+  """Helper for commands that take a build config in the arg list.
+
+  This function can cros_lib.Die() in some instances.
+
+  Args:
+    argv: A list of command line arguments.  If non-empty, [0] should be the
+        build spec.  These will not be modified.
+
+  Returns:
+    argv: The argv with the build spec stripped off.  This might be argv[1:] or
+        just argv.  Not guaranteed to be new memory.
+    build_config: The SafeConfigParser for the build config; might be None if
+        this is a host config.  TODO(dianders): Should there be a build spec for
+        the host?
+  """
+  # The spec name is optional.  If no arguments, we'll show a menu...
+  # Note that if there are arguments, but the first argument is a flag, we'll
+  # assume that we got called before OptionParser.  In that case, they might
+  # have specified options but still want the board menu.
+  if argv and not argv[0].startswith('-'):
+    spec_name = argv[0]
+    argv = argv[1:]
+  else:
+    spec_name = ''
+
+  spec_path = FindSpec(spec_name)
+
+  if spec_path == _HOST_TARGET:
+    return argv, None
+
+  build_config = ReadConfig(spec_path)
+
+  # TODO(dianders): Add a config checker class that makes sure that the
+  # target is not a blank string.  Might also be good to make sure that the
+  # target has no whitespace (so we don't screw up a subcommand invoked
+  # through a shell).
+
+  return argv, build_config
+
+
+def EnterChroot(chroot_config, fn, *args, **kwargs):
+  """Re-run the given function inside the chroot.
+
+  When the function is run, it will be run in a SEPARATE INSTANCE of chromite,
+  which will be run in the chroot.  This is a little weird.  Specifically:
+  - When the callee executes, it will be a separate python instance.
+    - Globals will be reset back to defaults.
+    - A different version of python (with different modules) may end up running
+      the script in the chroot.
+  - All arguments are pickled up into a temp file, whose path is passed on the
+    command line.
+    - That means that args must be pickleable.
+    - It also means that modifications to the parameters by the callee are not
+      visible to the caller.
+    - Even the function is "pickled".  The way the pickle works, I belive, is it
+      just passes the name of the function.  If this name somehow resolves
+      differently in the chroot, you may get weirdness.
+  - Since we're in the chroot, obviously files may have different paths.  It's
+    up to you to convert parameters if you need to.
+  - The stdin, stdout, and stderr aren't touched.
+
+  Args:
+    chroot_config: A SafeConfigParser representing the config for the chroot.
+    fn: Either: a) the function to call or b) A tuple of an object and the
+        name of the method to call.
+    args: All other arguments will be passed to the function as is.
+    kwargs: All other arguments will be passed to the function as is.
+  """
+  # Make sure that the chroot exists...
+  chroot_dir = GetChrootAbsDir(chroot_config)
+  if not DoesChrootExist(chroot_config):
+    cros_lib.Die(
+        'Chroot dir does not exist; try the "build host" command.\n  %s.' %
+        chroot_dir)
+
+  cros_lib.Info('ENTERING THE CHROOT')
+
+  # Save state to a temp file (inside the chroot!) using pickle.
+  tmp_dir = os.path.join(chroot_dir, 'tmp')
+  state_file = tempfile.NamedTemporaryFile(prefix='chromite', dir=tmp_dir)
+  try:
+    cPickle.dump((fn, args, kwargs), state_file, cPickle.HIGHEST_PROTOCOL)
+    state_file.flush()
+
+    # Translate temp file name into a chroot path...
+    chroot_state_path = os.path.join('/tmp', os.path.basename(state_file.name))
+
+    # Put together command.  We're going to force the shell to do all of the
+    # splitting of arguments, since we're throwing all of the flags from the
+    # config file in there.
+    # TODO(dianders): It might be nice to run chromite as:
+    #     python -m chromite.chromite_main
+    #   ...but, at the moment, that fails if you're in src/scripts
+    #   which already has a chromite folder.
+    cmd = (
+        './enter_chroot.sh --chroot="%s" %s --'
+        ' python ../../chromite/shell/main.py --resume-state %s') % (
+            chroot_dir,
+            chroot_config.get('CHROOT', 'enter_chroot_flags'),
+            chroot_state_path)
+
+    # We'll put CWD as src/scripts when running the command.  Since everyone
+    # running by hand has their cwd there, it is probably the safest.
+    cwd = os.path.join(SRCROOT_PATH, 'src', 'scripts')
+
+    # Run it.  We allow "error" so we don't print a confusing error message
+    # filled with out resume-state garbage on control-C.
+    cmd_result = cros_lib.RunCommand(cmd, shell=True, cwd=cwd, print_cmd=False,
+                            exit_code=True, error_ok=True, ignore_sigint=True)
+
+    if cmd_result.returncode:
+      cros_lib.Die('Chroot exited with error code %d' % cmd_result.returncode)
+  finally:
+    # Make sure things get closed (and deleted), even upon exception.
+    state_file.close()
+
+
+def ResumeEnterChrootIfNeeded(argv):
+  """Should be called as the first line in main() to support EnterChroot().
+
+  We'll check whether the --resume-state argument is there.  If the argument is
+  there, we'll use it and call the proper function (now that we're in the
+  chroot).  We'll then return True to indicate to main() that it should exit
+  immediately.
+
+  If the --resume-state argument is not there, this function will return False
+  without doing anything else.
+
+  Args:
+    argv: The value of sys.argv
+
+  Returns:
+    True if we resumed; indicates that main should return without doing any
+    further work.
+  """
+  if argv[1:2] == ['--resume-state']:
+    # Internal mechanism (not documented to users) to resume in the chroot.
+    # ...actual resume state file is passed in argv[2] for simplicity...
+    assert len(argv) == 3, 'Resume State not passed properly.'
+    fn, args, kwargs = cPickle.load(open(argv[2], 'rb'))
+
+    # Handle calling a method in a class; that can't be directly pickled.
+    if isinstance(fn, tuple):
+      obj, method = fn
+      fn = getattr(obj, method)
+
+    fn(*args, **kwargs)
+
+    # Return True to tell main() that it should exit.
+    return True
+  else:
+    # Return False to tell main() that we didn't find the --resume-state
+    # argument and that it should do normal arugment parsing.
+    return False