src/third_party/chromite/scripts/wrapper.py

   1 #!/usr/bin/python
   2 # Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
   3 # Use of this source code is governed by a BSD-style license that can be
   4 # found in the LICENSE file.
   5
   6 """Wrapper around chromite executable scripts.
   7
   8 This takes care of creating a consistent environment for chromite scripts
   9 (like setting up import paths) so we don't have to duplicate the logic in
  10 lots of places.
  11 """
  12
  13 from __future__ import print_function
  14
  15 import os
  16 import sys
  17
  18 CHROMITE_PATH = None
  19
  20
  21 class ChromiteImporter(object):
  22   """Virtual chromite module
  23
  24   If the checkout is not named 'chromite', trying to do 'from chromite.xxx'
  25   to import modules fails horribly.  Instead, manually locate the chromite
  26   directory (whatever it is named), load & return it whenever someone tries
  27   to import it.  This lets us use the stable name 'chromite' regardless of
  28   how things are structured on disk.
  29
  30   This also lets us keep the sys.path search clean.  Otherwise we'd have to
  31   worry about what other dirs chromite were checked out near to as doing an
  32   import would also search those for .py modules.
  33   """
  34
  35   # When trying to load the chromite dir from disk, we'll get called again,
  36   # so make sure to disable our logic to avoid an infinite loop.
  37   _loading = False
  38
  39   def find_module(self, fullname, _path):
  40     """Handle the 'chromite' module"""
  41     if fullname == 'chromite' and not self._loading:
  42       return self
  43     return None
  44
  45   def load_module(self, _fullname):
  46     """Return our cache of the 'chromite' module"""
  47     # Locate the top of the chromite dir by searching for the PRESUBMIT.cfg
  48     # file.  This assumes that file isn't found elsewhere in the tree.
  49     path = os.path.dirname(os.path.realpath(__file__))
  50     while not os.path.exists(os.path.join(path, 'PRESUBMIT.cfg')):
  51       path = os.path.dirname(path)
  52
  53     # pylint: disable=W0603
  54     global CHROMITE_PATH
  55     CHROMITE_PATH = path + '/'
  56
  57     # Finally load the chromite dir.
  58     path, mod = os.path.split(path)
  59     sys.path.insert(0, path)
  60     self._loading = True
  61     try:
  62       # This violates PEP302 slightly because __import__ will return the
  63       # cached module from sys.modules rather than reloading it from disk.
  64       # But the imp module does not work cleanly with meta_path currently
  65       # which makes it hard to use.  Until that is fixed, we won't bother
  66       # trying to address the edge case since it doesn't matter to us.
  67       return __import__(mod)
  68     finally:
  69       # We can't pop by index as the import might have changed sys.path.
  70       sys.path.remove(path)
  71       self._loading = False
  72
  73 sys.meta_path.insert(0, ChromiteImporter())
  74
  75 from chromite.lib import commandline
  76
  77
  78 def FindTarget(target):
  79   """Turn the path into something we can import from the chromite tree.
  80
  81   This supports a variety of ways of running chromite programs:
  82   # Loaded via depot_tools in $PATH.
  83   $ cros_sdk --help
  84   # Loaded via .../chromite/bin in $PATH.
  85   $ cros --help
  86   # No $PATH needed.
  87   $ ./bin/cros --help
  88   # Loaded via ~/bin in $PATH to chromite bin/ subdir.
  89   $ ln -s $PWD/bin/cros ~/bin; cros --help
  90   # No $PATH needed.
  91   $ ./buildbot/cbuildbot_config --help
  92   # No $PATH needed, but symlink inside of chromite dir.
  93   $ ln -s ./buildbot/cbuildbot_config; ./cbuildbot_config --help
  94   # Loaded via ~/bin in $PATH to non-chromite bin/ subdir.
  95   $ ln -s $PWD/buildbot/cbuildbot_config ~/bin/; cbuildbot_config --help
  96   # No $PATH needed, but a relative symlink to a symlink to the chromite dir.
  97   $ cd ~; ln -s bin/cbuildbot_config ./; ./cbuildbot_config --help
  98
  99   Args:
 100     target: Path to the script we're trying to run.
 101
 102   Returns:
 103     The module main functor.
 104   """
 105   while True:
 106     # Walk back one symlink at a time until we get into the chromite dir.
 107     parent, base = os.path.split(target)
 108     parent = os.path.realpath(parent)
 109     if parent.startswith(CHROMITE_PATH):
 110       target = base
 111       break
 112     target = os.path.join(os.path.dirname(target), os.readlink(target))
 113   assert parent.startswith(CHROMITE_PATH), (
 114       'could not figure out leading path\n'
 115       '\tparent: %s\n'
 116       '\tCHROMITE_PATH: %s' % (parent, CHROMITE_PATH))
 117   parent = parent[len(CHROMITE_PATH):].split(os.sep)
 118   target = ['chromite'] + parent + [target]
 119   # Our bin dir is just scripts stuff.
 120   if target[1] == 'bin':
 121     target[1] = 'scripts'
 122
 123   module = __import__('.'.join(target))
 124   # __import__ gets us the root of the namespace import; walk our way up.
 125   for attr in target[1:]:
 126     module = getattr(module, attr)
 127
 128   return getattr(module, 'main', None)
 129
 130
 131 if __name__ == '__main__':
 132   commandline.ScriptWrapperMain(FindTarget)