Complete reimplementation

- Work entirely in-place, even to switch to a different version of Python.
- Update bin/buildout to immediately use the wanted Python on subsequent
  buildout runs.

For SlapOS, the second point is required to have the instanciation done
with the built Python.

slapos/rebootstrap/__init__.py

 ##############################################################################
 #
-# Copyright (c) 2010 ViFiB SARL and Contributors.
+# Copyright (c) 2010-2017 ViFiB SARL and Contributors.
 # All Rights Reserved.
 #
 # This software is subject to the provisions of the Zope Public License,
@@ -11,52 +11,48 @@
 # FOR A PARTICULAR PURPOSE.
 #
 ##############################################################################
-import os
-import zc.buildout
-import zc.buildout.buildout
-import sys
-import logging
-import subprocess
 
-def extension(buildout):
-  Rebootstrap(buildout)()
+import logging, os, shutil, subprocess, sys, tempfile
+from zc.buildout import easy_install, UserError
+
+
+class extension(object):
 
-class Rebootstrap:
   def __init__(self, buildout):
-    self.logger = logging.getLogger(__name__)
     self.buildout = buildout
-    buildout_directory = buildout['buildout']['directory']
-    # fetch section to build python, obligatory
-    self.python_section = buildout['buildout'].get('python','').strip()
-    if not self.python_section:
-      raise zc.buildout.UserError('buildout:python is not defined.')
-    if self.python_section not in buildout:
-      raise zc.buildout.UserError('[%s] is not defined.' % self.python_section)
+    # fetch section to build python (default value is 'buildout')
+    self.python_section = buildout['buildout']['python'].strip()
     self.wanted_python = buildout[self.python_section]['executable']
-    rebootstrap_directory = buildout['buildout'].get('rebootstrap-directory')
-    if rebootstrap_directory:
-      self.rebootstrap_directory = os.path.join(
-        buildout_directory, 'rebootstrap'
-      )
-      self.wanted_python = self.wanted_python.replace(
-        buildout_directory, self.rebootstrap_directory, 1
-      )
-    else:
-      self.rebootstrap_directory = buildout_directory
-    # query for currently running python
-    self.running_python = sys.executable
-
-  def __call__(self):
-    if self.running_python != self.wanted_python:
-      self.install_section()
-      self.reboot()
-    elif self.python_section:
-      buildout = self.buildout['buildout']
-      if self.python_section not in buildout['parts'].split():
-        buildout['parts'] = self.python_section + '\n' + buildout['parts']
-
-  def reboot(self):
-    message = """
+    if sys.executable != self.wanted_python:
+      self.hook('_setup_directories')
+    elif buildout._parts:
+      self._frozen = frozenset(buildout._parts)
+      self.hook('_compute_part_signatures')
+
+  def hook(self, attr):
+    buildout = self.buildout
+    getattr(buildout, attr)
+    def wrapper(*args, **kw):
+      delattr(buildout, attr)
+      return getattr(self, attr)(*args, **kw)
+    setattr(buildout, attr, wrapper)
+
+  def _setup_directories(self):
+    logger = logging.getLogger(__name__)
+    buildout = self.buildout
+    logger.info(
+      "Make sure that the section %r won't be reinstalled after rebootstrap."
+      % self.python_section)
+    # We hooked in such a way that all extensions are loaded. Do not reload.
+    buildout._load_extensions
+    buildout._load_extensions = lambda: None
+    # workaround for the install command,
+    # which ignores dependencies when parts are specified
+    # (the only sections we have accessed so far are those that are required
+    #  to build the wanted Python)
+    buildout.install(buildout._parts) # [self.python_section]
+
+    logger.info("""
 ************ REBOOTSTRAP: IMPORTANT NOTICE ************
 bin/buildout is being reinstalled right now, as new python:
   %(wanted_python)s
@@ -64,90 +60,69 @@ is available, and buildout is using another python:
   %(running_python)s
 Buildout will be restarted automatically to have this change applied.
 ************ REBOOTSTRAP: IMPORTANT NOTICE ************
-""" % dict(wanted_python=self.wanted_python,
-    running_python=self.running_python)
-    self.logger.info(message)
-    args = sys.argv[:]
-    env = os.environ
-    if 'ORIG_PYTHON' not in env:
-      env['ORIG_PYTHON'] = sys.executable
-    os.execve(self.wanted_python, [self.wanted_python] + args, env)
-
-  def install_section(self):
-    if not os.path.exists(self.wanted_python) or \
-        self.rebootstrap_directory != self.buildout['buildout']['directory']:
-      self.logger.info('Installing section %r to provide %r' % (
-        self.python_section, self.wanted_python))
-      args = sys.argv[:]
-      if 'install' in args:
-        args =  args[:args.index('install')]
-
-      # explicitly specify the config file location by absolute path
-      if '-c' not in args:
-        config_file = os.path.abspath(os.path.join(
-          os.curdir, 'buildout.cfg'))
-        args.extend(['-c', config_file])
+""" % dict(wanted_python=self.wanted_python, running_python=sys.executable))
+
+    installed = sys.argv[0]
+    new_bin = installed + '-rebootstrap'
+    if not os.path.exists(new_bin):
+      from .bootstrap import get_distributions, setup_script
+      if subprocess.call((self.wanted_python, '-c',
+          'import sys; sys.exit(sys.version_info[:2] == %r)'
+          % (sys.version_info[:2],))):
+        setup_script(new_bin, self.wanted_python)
       else:
-        config_file = args[args.index('-c') + 1]
-        if not zc.buildout.buildout._isurl(config_file):
-          config_file = os.path.abspath(config_file)
-          args[args.index('-c') + 1] = config_file
-
-      # explicitly invoke with the current python interpreter
-      args.insert(0, sys.executable)
-
-      # remove rebootstrap extension, which is not needed in rebootstrap part
-      extension_list = self.buildout['buildout']['extensions'].split()
-      extension_list = [q.strip() for q in extension_list if q.strip() != \
-          __name__]
-      bin_directory = self.buildout['buildout']['bin-directory']
-      # rerun buildout with only neeeded section to reuse buildout
-      # ability to calcuate all dependency
-      args.extend([
-        # chroot to rebootstrap directory
-        'buildout:directory=%s' % self.rebootstrap_directory,
-        # preserve bin-directory outside the chroot.
-        'buildout:bin-directory=%s' % bin_directory,
-        # install only required section with dependencies
-        'buildout:parts=%s' % self.python_section,
-        # do not load this extension
-        'buildout:extensions=%s' % ' '.join(extension_list),
-        # more parameters for building slapos package
-        'buildout:rootdir=%s' % self.rebootstrap_directory,
-        'buildout:destdir=',
-      ])
-      self.logger.info('Rerunning buildout to install section %s with '
-        'arguments %r.'% (self.python_section, args))
-      process = subprocess.Popen(args)
-      process.wait()
-      if process.returncode != 0:
-        raise zc.buildout.UserError('Error during installing python '
-          'provision section.')
-    if not os.path.exists(self.wanted_python):
-      raise zc.buildout.UserError('Section %r directed to python executable:\n'
-          '%r\nUnfortunately even after installing this section executable was'
-          ' not found.\nThis is section responsibility to provide python (eg. '
-          'by compiling it).' % (self.python_section, self.wanted_python))
-
-_uninstall_part_orig = zc.buildout.buildout.Buildout._uninstall_part
-def _uninstall_part(self, part, installed_part_options):
-  _uninstall_part_orig(self, part, installed_part_options)
-  try:
-    location = self[part].get('location')
-  except zc.buildout.buildout.MissingSection:
-    return
-  if location and sys.executable.startswith(location):
-    message = """
-************ REBOOTSTRAP: IMPORTANT NOTICE ************
-%r part that provides the running Python is uninstalled.
-Buildout will be restarted automatically with the original Python.
-************ REBOOTSTRAP: IMPORTANT NOTICE ************
-""" % part
-    self._logger.info(message)
-    if getattr(self, 'dry_run', False):
-      sys.exit()
-    args = sys.argv[:]
-    env = os.environ
-    orig_python = env['ORIG_PYTHON']
-    os.execve(orig_python, [orig_python] + args, env)
-zc.buildout.buildout.Buildout._uninstall_part = _uninstall_part
+        # With a different version of Python,
+        # we must reinstall required eggs from source.
+        from pkg_resources import resource_string
+        with Cache(buildout['buildout']) as cache:
+          subprocess.check_call([self.wanted_python, '-c',
+            resource_string(__name__, 'bootstrap.py'),
+            new_bin, cache._dest, cache.tmp,
+            ] + list(map(cache, get_distributions())))
+
+    shutil.copy(new_bin, installed)
+    os.execv(self.wanted_python, [self.wanted_python] + sys.argv)
+
+  def _compute_part_signatures(self, install_parts):
+    # Skip signature check for parts that were required to build the wanted
+    # Python. Signatures differ when switching to a different version.
+    buildout = self.buildout
+    buildout._compute_part_signatures(install_parts)
+    installed_part_options = buildout.installed_part_options
+    for part in self._frozen.intersection(install_parts):
+      buildout[part]['__buildout_signature__'] = \
+        installed_part_options[part]['__buildout_signature__']
+
+
+class Cache(easy_install.Installer):
+
+  def __init__(self, buildout):
+    easy_install.Installer.__init__(self,
+      buildout['eggs-directory'],
+      buildout.get('find-links', '').split())
+
+  def __enter__(self):
+    self.tmp = self._download_cache or tempfile.mkdtemp('get_dist')
+    return self
+
+  def __exit__(self, t, v, tb):
+    if self.tmp is not self._download_cache:
+      shutil.rmtree(self.tmp)
+    del self.tmp
+
+  def __call__(self, dist):
+    req = dist.as_requirement()
+    cache = self._download_cache
+    if cache:
+      from pkg_resources import SOURCE_DIST
+      for avail in self._index[dist.project_name]:
+        if (avail.version == dist.version and
+            avail.precedence == SOURCE_DIST and
+            cache == os.path.dirname(avail.location)):
+          return str(req)
+    avail = self._obtain(req, True)
+    if avail is None:
+      raise UserError("Couldn't find a distribution for %r" % str(req))
+    if self._fetch(avail, self.tmp, cache) is None:
+      raise UserError("Couldn't download distribution %s." % avail)
+    return str(req)

slapos/rebootstrap/bootstrap.py

+import os, sys
+
+class FakeSysExecutable(object):
+
+  def __init__(self, python):
+    self.executable = python
+
+  def __getattr__(self, attr):
+    return getattr(sys, attr)
+
+def get_distributions():
+  from pkg_resources import get_distribution
+  distributions = ['setuptools', 'zc.buildout']
+  try:
+    import slapos.libnetworkcache
+  except ImportError:
+    pass
+  else:
+    distributions.append('slapos.libnetworkcache')
+  return map(get_distribution, distributions)
+
+def setup_script(path, python=sys.executable):
+  from zc.buildout import easy_install
+  try:
+    if sys.executable != python:
+      easy_install.sys = FakeSysExecutable(python)
+    easy_install.scripts(
+      ((os.path.realpath(path), 'zc.buildout.buildout', 'main'),),
+      get_distributions(),
+      python)
+  finally:
+    easy_install.sys = sys
+
+def main():
+  import shutil, subprocess, tempfile, zipfile
+  eggs_dir = sys.argv[2]
+  cache = sys.argv[3]
+  # Install setuptools.
+  src = os.path.join(cache, sys.argv[4].replace('==', '-') + '.zip')
+  tmp = tempfile.mkdtemp()
+  try:
+    with zipfile.ZipFile(src) as zip_file:
+      zip_file.extractall(tmp)
+    src, = os.listdir(tmp)
+    subprocess.check_call((sys.executable, 'setup.py', '-q', 'bdist_egg',
+                          '--dist-dir', tmp), cwd=os.path.join(tmp, src))
+    egg = os.listdir(tmp)
+    egg.remove(src)
+    egg, = egg
+    dst = os.path.join(eggs_dir, egg)
+    os.path.exists(dst) or shutil.move(os.path.join(tmp, egg), dst)
+  finally:
+    shutil.rmtree(tmp)
+  sys.path.insert(0, dst)
+  # Install other requirements given on command line.
+  from pkg_resources import working_set, require
+  from setuptools.command import easy_install
+  reqs = sys.argv[5:]
+  easy_install.main(['-mZqNxd', eggs_dir, '-f', cache] + reqs)
+  working_set.add_entry(eggs_dir)
+  for req in reqs:
+    require(req)
+  # Generate bin/buildout-rebootstrap script.
+  setup_script(sys.argv[1])
+
+if __name__ == '__main__':
+  sys.exit(main())