Adds Python package management to depot_tools.

package_management.py

Index: package_management.py
diff --git a/package_management.py b/package_management.py
new file mode 100644
index 0000000000000000000000000000000000000000..61e7ad28641c3aa8167f9b81110069c5be6775ff
--- /dev/null
+++ b/package_management.py
@@ -0,0 +1,602 @@
+#!/usr/bin/env python
+# Copyright (c) 2012 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Package Management
+
+This module is used to bring in external Python dependencies via eggs from the
+PyPi repository.
+
+The approach is to create a site directory in depot_tools root which will
+contain those packages that are listed as dependencies in the module variable
+PACKAGES. Since we can't guarantee that setuptools is available in all
+distributions this module also contains the ability to bootstrap the site
+directory by manually downloading and installing setuptools. Once setuptools is
+available it uses that to install the other packages in the traditional
+manner.
+
+Use is simple:
+
+  import package_management
+
+  # Before any imports from the site directory, call this. This only needs
+  # to be called in one place near the beginning of the program.
+  package_management.SetupSiteDirectory()
+
+  # If 'SetupSiteDirectory' fails it will complain with an error message but
+  # continue happily. Expect ImportErrors when trying to import any third
+  # party modules from the site directory.
+
+  import some_third_party_module
+
+  ... etc ...
+"""
+
+import cStringIO
+import os
+import re
+import shutil
+import site
+import subprocess
+import sys
+import tempfile
+import urllib2
+
+
+# This is the version of setuptools that we will download if the local
+# python distribution does not include one.
+SETUPTOOLS = ('setuptools', '0.6c11')
+
+# These are the packages that are to be installed in the site directory.
+# easy_install makes it so that the most recently installed version of a
+# package is the one that takes precedence, even if a newer version exists
+# in the site directory. This allows us to blindly install these one on top
+# of the other without worrying about whats already installed.
+#
+# NOTE: If we are often rolling these dependencies then users' site
+#     directories will grow monotonically. We could be purging any orphaned
+#     packages using the tools provided by pkg_resources.
+PACKAGES = (('logilab-common', '0.57.1'),
+            ('logilab-astng', '0.23.1'),
+            ('pylint', '0.25.1'))
+
+
+# The Python version suffix used in generating the site directory and in
+# requesting packages from PyPi.
+VERSION_SUFFIX = "%d.%d" % sys.version_info[0:2]
+
+# This is the root directory of the depot_tools installation.
+ROOT_DIR = os.path.abspath(os.path.dirname(__file__))
+
+# This is the path of the site directory we will use. We make this
+# python version specific so that this will continue to work even if the
+# python version is rolled.
+SITE_DIR = os.path.join(ROOT_DIR, 'site-packages-py%s' % VERSION_SUFFIX)
+
+# This status file is created the last time PACKAGES were rolled successfully.
+# It is used to determine if packages need to be rolled by comparing against
+# the age of __file__.
+LAST_ROLLED = os.path.join(SITE_DIR, 'last_rolled.txt')
+
+
+class Error(Exception):
+  """The base class for all module errors."""
+  pass
+
+
+class InstallError(Error):
+  """Thrown if an installation is unable to complete."""
+  pass
+
+
+class Package(object):
+  """A package represents a release of a project.
+
+  We use this as a lightweight version of pkg_resources, allowing us to
+  perform an end-run around setuptools for the purpose of bootstrapping. Its
+  functionality is very limited.
+
+  Attributes:
+    name: the name of the package.
+    version: the version of the package.
+    safe_name: the safe name of the package.
+    safe_version: the safe version string of the package.
+    file_name: the filename-safe name of the package.
+    file_version: the filename-safe version string of the package.
+  """
+
+  def __init__(self, name, version):
+    """Initialize this package.
+
+    Args:
+      name: the name of the package.
+      version: the version of the package.
+    """
+    self.name = name
+    self.version = version
+    self.safe_name = Package._MakeSafeName(self.name)
+    self.safe_version = Package._MakeSafeVersion(self.version)
+    self.file_name = Package._MakeSafeForFilename(self.safe_name)
+    self.file_version = Package._MakeSafeForFilename(self.safe_version)
+
+  @staticmethod
+  def _MakeSafeName(name):
+    """Makes a safe package name, as per pkg_resources."""
+    return re.sub('[^A-Za-z0-9]+', '-', name)
+
+  @staticmethod
+  def _MakeSafeVersion(version):
+    """Makes a safe package version string, as per pkg_resources."""
+    version = re.sub('\s+', '.', version)
+    return re.sub('[^A-Za-z0-9\.]+', '-', version)
+
+  @staticmethod
+  def _MakeSafeForFilename(safe_name_or_version):
+    """Makes a safe name or safe version safe to use in a file name.
+    |safe_name_or_version| must be a safe name or version string as returned
+    by GetSafeName or GetSafeVersion.
+    """
+    return re.sub('-', '_', safe_name_or_version)
+
+  def GetAsRequirementString(self):
+    """Builds an easy_install requirements string representing this package."""
+    return '%s==%s' % (self.name, self.version)
+
+  def GetFilename(self, extension=None):
+    """Builds a filename for this package using the setuptools convention.
+    If |extensions| is provided it will be appended to the generated filename,
+    otherwise only the basename is returned.
+
+    The following url discusses the filename format:
+
+      http://svn.python.org/projects/sandbox/trunk/setuptools/doc/formats.txt
+    """
+    filename = '%s-%s-py%s' % (self.file_name, self.file_version,
+                               VERSION_SUFFIX)
+
+    if extension:
+      if not extension.startswith('.'):
+        filename += '.'
+      filename += extension
+
+    return filename
+
+  def GetPyPiUrl(self, extension):
+    """Returns the URL where this package is hosted on PyPi."""
+    return 'http://pypi.python.org/packages/%s/%c/%s/%s' % (VERSION_SUFFIX,
+        self.file_name[0], self.file_name, self.GetFilename(extension))
+
+  def DownloadEgg(self, dest_dir, overwrite=False):
+    """Downloads the EGG for this URL.
+
+    Args:
+      dest_dir: The directory where the EGG should be written. If the EGG
+          has already been downloaded and cached the returned path may not
+          be in this directory.
+      overwite: If True the destination path will be overwritten even if
+          it already exists. Defaults to False.
+
+    Returns:
+      The path to the written EGG.
+
+    Raises:
+      Error: if dest_dir doesn't exist, the EGG is unable to be written,
+          the URL doesn't exist, or the server returned an error, or the
+          transmission was interrupted.
+    """
+    if not os.path.exists(dest_dir):
+      raise Error('Path does not exist: %s' % dest_dir)
+
+    if not os.path.isdir(dest_dir):
+      raise Error('Path is not a directory: %s' % dest_dir)
+
+    filename = os.path.abspath(os.path.join(dest_dir, self.GetFilename('egg')))
+    if os.path.exists(filename):
+      if os.path.isdir(filename):
+        raise Error('Path is a directory: %s' % filename)
+      if not overwrite:
+        return filename
+
+    url = self.GetPyPiUrl('egg')
+
+    try:
+      url_stream = urllib2.urlopen(url)
+      local_file = open(filename, 'wb')
+      local_file.write(url_stream.read())
+      local_file.close()
+      return filename
+    except (IOError, urllib2.HTTPError, urllib2.URLError):
+      # Reraise with a new error type, keeping the original message and
+      # traceback.
+      raise Error, sys.exc_info()[1], sys.exc_info()[2]
+
+
+def AddToPythonPath(path):
+  """Adds the provided path to the head of PYTHONPATH and sys.path."""
+  path = os.path.abspath(path)
+  if path not in sys.path:
+    sys.path.insert(0, path)
+
+  paths = os.environ.get('PYTHONPATH', '').split(os.pathsep)
+  if path not in paths:
+    paths.insert(0, path)
+    os.environ['PYTHONPATH'] = os.pathsep.join(paths)
+
+
+def AddSiteDirectory(path):
+  """Adds the provided path to the runtime as a site directory.
+
+  Any modules that are in the site directory will be available for importing
+  after this returns. If modules are added or deleted this must be called
+  again for the changes to be reflected in the runtime.
+
+  This calls both AddToPythonPath and site.addsitedir. Both are needed to
+  convince easy_install to treat |path| as a site directory.
+  """
+  AddToPythonPath(path)
+  site.addsitedir(path)  # pylint: disable=E1101
+
+def EnsureSiteDirectory(path):
+  """Creates and/or adds the provided path to the runtime as a site directory.
+
+  This works like AddSiteDirectory but it will create the directory if it
+  does not yet exist.
+
+  Raise:
+    Error: if the site directory is unable to be created, or if it exists and
+        is not a directory.
+  """
+  if os.path.exists(path):
+    if not os.path.isdir(path):
+      raise Error('Path is not a directory: %s' % path)
+  else:
+    try:
+      os.mkdir(path)
+    except IOError:
+      raise Error('Unable to create directory: %s' % path)
+
+  AddSiteDirectory(path)
+
+
+def ModuleIsFromPackage(module, package_path):
+  """Determines if a module has been imported from a given package.
+
+  Args:
+    module: the module to test.
+    package_path: the path to the package to test.
+
+  Returns:
+    True if |module| has been imported from |package_path|, False otherwise.
+  """
+  try:
+    m = os.path.abspath(module.__file__)
+    p = os.path.abspath(package_path)
+    if len(m) <= len(p):
+      return False
+    if m[0:len(p)] != p:
+      return False
+    return m[len(p)] == os.sep
+  except AttributeError:
+    return False
+
+
+def _CaptureStdStreams(function, *args, **kwargs):
+  """Captures stdout and stderr while running the provided function.
+
+  This only works if |function| only accesses sys.stdout and sys.stderr. If
+  we need more than this we'll have to use subprocess.Popen.
+
+  Args:
+    function: the function to be called.
+    args: the arguments to pass to |function|.
+    kwargs: the keyword arguments to pass to |function|.
+  """
+  orig_stdout = sys.stdout
+  orig_stderr = sys.stderr
+  sys.stdout = cStringIO.StringIO()
+  sys.stderr = cStringIO.StringIO()
+  try:
+    return function(*args, **kwargs)
+  finally:
+    sys.stdout = orig_stdout
+    sys.stderr = orig_stderr
+
+
+def InstallPackage(url_or_req, site_dir):
+  """Installs a package to a site directory.
+
+  |site_dir| must exist and already be an active site directory. setuptools
+  must in the path. Uses easy_install which may involve a download from
+  pypi.python.org, so this also requires network access.
+
+  Args:
+    url_or_req: the package to install, expressed as an URL (may be local),
+        or a requirement string.
+    site_dir: the site directory in which to install it.
+
+  Raises:
+    InstallError: if installation fails for any reason.
+  """
+  args = ['--quiet', '--install-dir', site_dir, '--exclude-scripts',
+          '--always-unzip', '--no-deps', url_or_req]
+
+  # The easy_install script only calls SystemExit if something goes wrong.
+  # Otherwise, it falls through returning None.
+  try:
+    import setuptools.command.easy_install
+    _CaptureStdStreams(setuptools.command.easy_install.main, args)
+  except (ImportError, SystemExit):
+    # Re-raise the error, preserving the stack trace and message.
+    raise InstallError, sys.exc_info()[1], sys.exc_info()[2]
+
+
+def _RunInSubprocess(pycode):
+  """Launches a python subprocess with the provided code.
+
+  The subprocess will be launched with the same stdout and stderr. The
+  subprocess will use the same instance of python as is currently running,
+  passing |pycode| as arguments to this script. |pycode| will be interpreted
+  as python code in the context of this module.
+
+  Returns:
+    True if the subprocess returned 0, False if it returned an error.
+  """
+  return not subprocess.call([sys.executable, __file__, pycode])
+
+
+def _LoadSetupToolsFromEggAndInstall(egg_path):
+  """Loads setuptools from the provided egg |egg_path|, and installs it to
+  SITE_DIR.
+
+  This is intended to be run from a subprocess as it pollutes the running
+  instance of Python by importing a module and then forcibly deleting its
+  source.
+
+  Returns:
+    True on success, False on failure.
+  """
+  AddToPythonPath(egg_path)
+
+  try:
+    # Import setuptools and ensure it comes from the EGG.
+    import setuptools
+    if not ModuleIsFromPackage(setuptools, egg_path):
+      raise ImportError()
+  except ImportError:
+    print '  Unable to import downloaded package!'
+    return False
+
+  try:
+    print '  Using setuptools to install itself ...'
+    InstallPackage(egg_path, SITE_DIR)
+  except InstallError:
+    print '  Unable to install setuptools!'
+    return False
+
+  return True
+
+
+def BootstrapSetupTools():
+  """Bootstraps the runtime with setuptools.
+
+  Will try to import setuptools directly. If not found it will attempt to
+  download it and load it from there. If the download is successful it will
+  then use setuptools to install itself in the site directory.
+
+  This is meant to be run from a child process as it modifies the running
+  instance of Python by importing modules and then physically deleting them
+  from disk.
+
+  Returns:
+    Returns True if 'import setuptools' will succeed, False otherwise.
+  """
+  AddSiteDirectory(SITE_DIR)
+
+  # Check if setuptools is already available. If so, we're done.
+  try:
+    import setuptools  # pylint: disable=W0612
+    return True
+  except ImportError:
+    pass
+
+  print 'Bootstrapping setuptools ...'
+
+  EnsureSiteDirectory(SITE_DIR)
+
+  # Download the egg to a temp directory.
+  dest_dir = tempfile.mkdtemp('depot_tools')
+  path = None
+  try:
+    package = Package(*SETUPTOOLS)
+    print '  Downloading %s ...' % package.GetFilename()
+    path = package.DownloadEgg(dest_dir)
+  except Error:
+    print '  Download failed!'
+    shutil.rmtree(dest_dir)
+    return False
+
+  try:
+    # Load the downloaded egg, and install it to the site directory. Do this
+    # in a subprocess so as not to pollute this runtime.
+    pycode = '_LoadSetupToolsFromEggAndInstall(%s)' % repr(path)
+    if not _RunInSubprocess(pycode):
+      raise Error()
+
+    # Reload our site directory, which should now contain setuptools.
+    AddSiteDirectory(SITE_DIR)
+
+    # Try to import setuptools
+    import setuptools
+  except ImportError:
+    print '  Unable to import setuptools!'
+    return False
+  except Error:
+    # This happens if RunInSubProcess fails, and the appropriate error has
+    # already been written to stdout.
+    return False
+  finally:
+    # Delete the temp directory.
+    shutil.rmtree(dest_dir)
+
+  return True
+
+
+def _GetModTime(path):
+  """Gets the last modification time associated with |path| in seconds since
+  epoch, returning 0 if |path| does not exist.
+  """
+  try:
+    return os.stat(path).st_mtime
+  except:  # pylint: disable=W0702
+    # This error is different depending on the OS, hence no specified type.
+    return 0
+
+
+def _SiteDirectoryIsUpToDate():
+  return _GetModTime(LAST_ROLLED) > _GetModTime(__file__)
+
+
+def UpdateSiteDirectory():
+  """Installs the packages from PACKAGES if they are not already installed.
+  At this point we must have setuptools in the site directory.
+
+  This is intended to be run in a subprocess *prior* to the site directory
+  having been added to the parent process as it may cause packages to be
+  added and/or removed.
+
+  Returns:
+    True on success, False otherwise.
+  """
+  if _SiteDirectoryIsUpToDate():
+    return True
+
+  try:
+    AddSiteDirectory(SITE_DIR)
+    import pkg_resources
+
+    # Determine if any packages actually need installing.
+    missing_packages = []
+    for package in [SETUPTOOLS] + list(PACKAGES):
+      pkg = Package(*package)
+      req = pkg.GetAsRequirementString()
+
+      # It may be that this package is already available in the site
+      # directory. If so, we can skip past it without trying to install it.
+      pkg_req = pkg_resources.Requirement.parse(req)
+      try:
+        dist = pkg_resources.working_set.find(pkg_req)
+        if dist:
+          continue
+      except pkg_resources.VersionConflict:
+        # This happens if another version of the package is already
+        # installed in another site directory (ie: the system site directory).
+        pass
+
+      missing_packages.append(pkg)
+
+    # Install the missing packages.
+    if missing_packages:
+      print 'Updating python packages ...'
+      for pkg in missing_packages:
+        print '  Installing %s ...' % pkg.GetFilename()
+        InstallPackage(pkg.GetAsRequirementString(), SITE_DIR)
+
+    # Touch the status file so we know that we're up to date next time.
+    open(LAST_ROLLED, 'wb')
+  except InstallError, e:
+    print '  Installation failed: %s' % str(e)
+    return False
+
+  return True
+
+
+def SetupSiteDirectory():
+  """Sets up the site directory, bootstrapping setuptools if necessary.
+
+  If this finishes successfully then SITE_DIR will exist and will contain
+  the appropriate version of setuptools and all of the packages listed in
+  PACKAGES.
+
+  This is the main workhorse of this module. Calling this will do everything
+  necessary to ensure that you have the desired packages installed in the
+  site directory, and the site directory enabled in this process.
+
+  Returns:
+    True on success, False on failure.
+  """
+  if _SiteDirectoryIsUpToDate():
+    AddSiteDirectory(SITE_DIR)
+    return True
+
+  if not _RunInSubprocess('BootstrapSetupTools()'):
+    return False
+
+  if not _RunInSubprocess('UpdateSiteDirectory()'):
+    return False
+
+  # Process the site directory so that the packages within it are available
+  # for import.
+  AddSiteDirectory(SITE_DIR)
+
+  return True
+
+
+def CanImportFromSiteDirectory(package_name):
+  """Determines if the given package can be imported from the site directory.
+
+  Args:
+    package_name: the name of the package to import.
+
+  Returns:
+    True if 'import package_name' will succeed and return a module from the
+    site directory, False otherwise.
+  """
+  try:
+    return ModuleIsFromPackage(__import__(package_name), SITE_DIR)
+  except ImportError:
+    return False
+
+
+def Test():
+  """Runs SetupSiteDirectory and then tries to load pylint, ensuring that it
+  comes from the site directory just created. This is an end-to-end unittest
+  and allows for simple testing from the command-line by running
+
+    ./package_management.py 'Test()'
+  """
+  print 'Testing package_management.'
+  if not SetupSiteDirectory():
+    print 'SetupSiteDirectory failed.'
+    return False
+  if not CanImportFromSiteDirectory('pylint'):
+    print 'CanImportFromSiteDirectory failed.'
+    return False
+  print 'Success!'
+  return True
+
+
+def Main():
+  """The main entry for the package management script.
+
+  If no arguments are provided simply runs SetupSiteDirectory. If arguments
+  have been passed we execute the first argument as python code in the
+  context of this module. This mechanism is used during the bootstrap
+  process so that the main instance of Python does not have its runtime
+  polluted by various intermediate packages and imports.
+
+  Returns:
+    0 on success, 1 otherwise.
+  """
+  if len(sys.argv) == 2:
+    result = False
+    exec('result = %s' % sys.argv[1])
+
+    # Translate the success state to a return code.
+    return not result
+  else:
+    return not SetupSiteDirectory()
+
+
+if __name__ == '__main__':
+  sys.exit(Main())