Allow CROS_WORKON_SRCROOT to be used outside the chroot.

bin/chromite

 #!/usr/bin/python
 # 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.
 
 """Wrapper for the chromite shell.
 
 This script is intended to run in several ways:
 - Outside the chroot, it should be _copied_ to someplace that is in the
   path (like depot_tools).  It will search for the right chromite by looking
   for a file 'chromite/shell/main.py' upward based on the CWD.
 - Inside the chroot, it might be _either_ copied to someplace in the path (since
   depot_tools is in the path in the chroot) or it might run from chromite/bin
   directly, which should be in the PATH.  In any case, we'll look for the
   real 'chromite/shell/main.py' based on the environment variable
   CROS_WORKON_SRCROOT, so it doesn't matter what the CWD is.
 
 If you're looking at a copy and want to know where the original looks at, look
 here:
   http://git.chromium.org/gitweb/?p=chromite.git;a=blob;f=bin/chromite
 
 Since this script is _copied_, it should remain small and not use internal libs.
 """
 
 # Python imports.
 import os
 import sys
 
+
+def ReportChromiteError(message):

[diandersAtChromium, 2011/02/02 23:15:28]

nit: Should probably start w/ _, since this is private.  Should also have
docstring.

+  print >>sys.stderr, (
+    "ERROR: Couldn't find the chromite tool.\n\n" + message +
+    "\n\nFor details and help, see:\n"
+    "  http://www.chromium.org/chromium-os/developer-guide"
+  )
+
+
+def TryChromiteShell(dir_):

[diandersAtChromium, 2011/02/02 23:15:28]

nit: Should probably start w/ _, since this is private.  Should also have
docstring.

+  sys.path[1] = dir_
+  try:
+    import chromite.shell.main
+    return chromite.shell.main.main
+  except ImportError:
+    return None
+
+
 if __name__ == '__main__':
-  # Look relative to CROS_WORKON_SRCROOT if that variable exists.  This is
-  # the "inside the chroot" case.
+  sys.path.insert(1, None)

[diandersAtChromium, 2011/02/02 23:15:28]

A few issues here:

1. We don't need to insert anything into sys.path if we're inside the chroot. 
Don't insert anything in that case.

2. This line needs some comments, since it's not obvious.  Why inserting None? 
Why inserting at position 1 (as opposed to position 0)?

+  # Look relative to CROS_WORKON_SRCROOT if that variable exists.
+  # Inside the chroot, this variable should always be set; outside
+  # the setting is optional.
   if 'CROS_WORKON_SRCROOT' in os.environ:
-    chromite_path = os.path.join(os.environ['CROS_WORKON_SRCROOT'],
-                                 'chromite', 'shell', 'main.py')
-    if os.path.isfile(chromite_path):
-      # Exec the script, which will never return.
-      os.execv(chromite_path, sys.argv)
+    main_function = TryChromiteShell(os.environ['CROS_WORKON_SRCROOT'])
+    if main_function is None:
+      ReportChromiteError(
+        "There is no chromite package in the path indicated by\n"
+        "CROS_WORKON_SRCROOT.  The variable may be set wrong, or you\n"
+        "may need to update your chroot.")
+      sys.exit(1)
+  else:
+    # We're outside the chroot and CROS_WORKON_SRCROOT isn't set.
+    # Search upward until the parent dir doesn't change (on Linux,
+    # that means we're at '/').
+    dir_ = os.getcwd()
+    prev_dir = None
+    while dir_ != prev_dir:
+      main_function = TryChromiteShell(dir_)
+      if main_function is not None:
+        break
+      prev_dir = dir_
+      dir_ = os.path.dirname(dir_)
     else:
-      print (
-          "ERROR: Couldn't find the chromite tool.\n"
-          "\n"
-          "You may need to update your chroot.  If you need help, see:\n"
-          "  http://www.chromium.org/chromium-os/developer-guide"
-      )
+      ReportChromiteError(
+          "Please change to a directory inside your Chromium OS source tree\n"
+          "and retry.")
       sys.exit(1)
 
-  # Outside the chroot, search upward until the "parent" dir doesn't change
-  # (on Linux, that means we're at '/').  That is an error case.
-  dir = os.getcwd()
-  prev_dir = None
-  while dir != prev_dir:
-    chromite_path = os.path.join(dir, 'chromite', 'shell', 'main.py')
-    if os.path.isfile(chromite_path):
-      # Add the directory above chromite to PYTHONPATH before executing, so
-      # that "import chromite.abc.xyz" works...
-      env = os.environ.copy()
-      if 'PYTHONPATH' in env:
-        env['PYTHONPATH'] += ':%s' % dir
-      else:
-        env['PYTHONPATH'] = dir
-
-      # Exec the script, which will never return.
-      os.execve(chromite_path, sys.argv, env)
-
-    prev_dir = dir
-    dir = os.path.dirname(dir)
-
-  # TODO(dianders): Should we actually print out the 'repo init' call that
-  # the user should use?
-  print (
-      "ERROR: Couldn't find the chromite tool.\n"
-      "\n"
-      "Please change to a directory inside your Chromium OS source tree\n"
-      "and retry.  If you need to setup a Chromium OS source tree, see:\n"
-      "  http://www.chromium.org/chromium-os/developer-guide"
-  )
-  sys.exit(1)
+main_function()

[diandersAtChromium, 2011/02/02 23:15:28]

Please indent by 2 spaces so this is in the __name__ == '__main__' case.

+sys.exit(0)

[diandersAtChromium, 2011/02/02 23:15:28]

Why is sys.exit needed?