Fixup pex re-exec during bootstrap.

pex/interpreter.py

@@ -13,6 +13,7 @@
 
 from pex.compatibility import string
 from pex.executor import Executor
+from pex.orderedset import OrderedSet
 from pex.pep425tags import (
     get_abbr_impl,
     get_abi_tag,
@@ -283,7 +284,7 @@ def get(cls):
   @classmethod
   def all(cls, paths=None):
     if paths is None:
-      paths = os.getenv('PATH', '').split(':')
+      paths = os.getenv('PATH', '').split(os.pathsep)
     return cls.filter(cls.find(paths))
 
   @classmethod
@@ -334,12 +335,13 @@ def from_binary(cls, binary):
              extras.
     :rtype: :class:`PythonInterpreter`
     """
-    if binary not in cls.CACHE:
-      if binary == sys.executable:
-        cls.CACHE[binary] = cls._from_binary_internal()
+    normalized_binary = os.path.realpath(binary)
+    if normalized_binary not in cls.CACHE:
+      if normalized_binary == os.path.realpath(sys.executable):
+        cls.CACHE[normalized_binary] = cls._from_binary_internal()
       else:
-        cls.CACHE[binary] = cls._from_binary_external(binary)
-    return cls.CACHE[binary]
+        cls.CACHE[normalized_binary] = cls._from_binary_external(normalized_binary)
+    return cls.CACHE[normalized_binary]
 
   @classmethod
   def _matches_binary_name(cls, basefile):
@@ -378,7 +380,7 @@ def version_filter(version):
       return (version[MAJOR] == 2 and version[MINOR] >= 7 or
               version[MAJOR] == 3 and version[MINOR] >= 4)
 
-    all_versions = set(interpreter.identity.version for interpreter in pythons)
+    all_versions = OrderedSet(interpreter.identity.version for interpreter in pythons)
     good_versions = filter(version_filter, all_versions)
 
     for version in good_versions:

pex/pex_bootstrapper.py

@@ -18,122 +18,159 @@
 __all__ = ('bootstrap_pex',)
 
 
-def find_in_path(target_interpreter):
-  if os.path.exists(target_interpreter):
-    return target_interpreter
-
-  for directory in os.getenv('PATH', '').split(os.pathsep):
-    try_path = os.path.join(directory, target_interpreter)
-    if os.path.exists(try_path):
-      return try_path
+def _find_pex_python(pex_python):
+  def try_create(try_path):
+    try:
+      return PythonInterpreter.from_binary(try_path)
+    except Executor.ExecutionError:
+      return None
+
+  interpreter = try_create(pex_python)
+  if interpreter:
+    # If the target interpreter specified in PEX_PYTHON is an existing absolute path - use it.
+    yield interpreter
+  else:
+    # Otherwise scan the PATH for matches:
+    for directory in os.getenv('PATH', '').split(os.pathsep):
+      try_path = os.path.join(directory, pex_python)
+      interpreter = try_create(try_path)
+      if interpreter:
+        yield interpreter
 
 
-def find_compatible_interpreters(pex_python_path=None, compatibility_constraints=None):
+def find_compatible_interpreters(path=None, compatibility_constraints=None):
   """Find all compatible interpreters on the system within the supplied constraints and use
-     PEX_PYTHON_PATH if it is set. If not, fall back to interpreters on $PATH.
+     path if it is set. If not, fall back to interpreters on $PATH.
   """
-  if pex_python_path:
-    interpreters = []
-    for binary in pex_python_path.split(os.pathsep):
-      try:
-        interpreters.append(PythonInterpreter.from_binary(binary))
-      except Executor.ExecutionError:
-        print("Python interpreter %s in PEX_PYTHON_PATH failed to load properly." % binary,
-          file=sys.stderr)
-    if not interpreters:
-      die('PEX_PYTHON_PATH was defined, but no valid interpreters could be identified. Exiting.')
+  interpreters = OrderedSet()
+  paths = None
+  if path:
+    paths = path.split(os.pathsep)
   else:
-    # We may have been invoked with a specific interpreter not on the $PATH, make sure our
-    # sys.executable is included as a candidate in this case.
-    interpreters = OrderedSet([PythonInterpreter.get()])
+    # We may have been invoked with a specific interpreter, make sure our sys.executable is included
+    # as a candidate in this case.
+    interpreters.add(PythonInterpreter.get())
+  interpreters.update(PythonInterpreter.all(paths=paths))
+  return _filter_compatible_interpreters(interpreters,
+                                         compatibility_constraints=compatibility_constraints)
 
-    # Add all qualifying interpreters found in $PATH.
-    interpreters.update(PythonInterpreter.all())
 
+def _filter_compatible_interpreters(interpreters, compatibility_constraints=None):
   return list(
     matched_interpreters(interpreters, compatibility_constraints)
     if compatibility_constraints
     else interpreters
   )
 
 
-def _select_pex_python_interpreter(target_python, compatibility_constraints=None):
-  target = find_in_path(target_python)
-
-  if not target:
-    die('Failed to find interpreter specified by PEX_PYTHON: %s' % target)
-  if compatibility_constraints:
-    pi = PythonInterpreter.from_binary(target)
-    if not list(matched_interpreters([pi], compatibility_constraints)):
-      die('Interpreter specified by PEX_PYTHON (%s) is not compatible with specified '
-          'interpreter constraints: %s' % (target, str(compatibility_constraints)))
-  if not os.path.exists(target):
-    die('Target interpreter specified by PEX_PYTHON %s does not exist. Exiting.' % target)
-  return target
+def _select_pex_python_interpreter(pex_python, compatibility_constraints=None):
+  compatible_interpreters = _filter_compatible_interpreters(
+    _find_pex_python(pex_python),
+    compatibility_constraints=compatibility_constraints
+  )
+  if not compatible_interpreters:
+    die('Failed to find a compatible PEX_PYTHON={} for constraints: {}'
+        .format(pex_python, compatibility_constraints))
+  return _select_interpreter(compatible_interpreters)
 
 
-def _select_interpreter(pex_python_path=None, compatibility_constraints=None):
+def _select_path_interpreter(path=None, compatibility_constraints=None):
   compatible_interpreters = find_compatible_interpreters(
-    pex_python_path=pex_python_path, compatibility_constraints=compatibility_constraints)
-
+    path=path,
+    compatibility_constraints=compatibility_constraints
+  )
   if not compatible_interpreters:
-    die('Failed to find compatible interpreter for constraints: %s'
-        % str(compatibility_constraints))
-  # TODO: https://github.com/pantsbuild/pex/issues/430
-  target = min(compatible_interpreters).binary
+    die('Failed to find compatible interpreter on path {} for constraints: {}'
+        .format(path or os.getenv('PATH'), compatibility_constraints))
+  return _select_interpreter(compatible_interpreters)
 
-  if os.path.exists(target):
-    return target
 
+def _select_interpreter(candidate_interpreters):
+  current_interpreter = PythonInterpreter.get()
+  if current_interpreter in candidate_interpreters:
+    # Always prefer continuing with the current interpreter when possible.
+    return current_interpreter
+  else:
+    # TODO: Allow the selection strategy to be parameterized:
+    #   https://github.com/pantsbuild/pex/issues/430
+    return min(candidate_interpreters)
 
-def maybe_reexec_pex(compatibility_constraints):
-  """
-  Handle environment overrides for the Python interpreter to use when executing this pex.
+
+def maybe_reexec_pex(compatibility_constraints=None):
+  """Handle environment overrides for the Python interpreter to use when executing this pex.
 
   This function supports interpreter filtering based on interpreter constraints stored in PEX-INFO
-  metadata. If PEX_PYTHON is set in a pexrc, it attempts to obtain the binary location of the
-  interpreter specified by PEX_PYTHON. If PEX_PYTHON_PATH is set, it attempts to search the path for
-  a matching interpreter in accordance with the interpreter constraints. If both variables are
-  present in a pexrc, this function gives precedence to PEX_PYTHON_PATH and errors out if no
-  compatible interpreters can be found on said path.
+  metadata. If PEX_PYTHON is set it attempts to obtain the binary location of the interpreter
+  specified by PEX_PYTHON. If PEX_PYTHON_PATH is set, it attempts to search the path for a matching
+  interpreter in accordance with the interpreter constraints. If both variables are present, this
+  function gives precedence to PEX_PYTHON_PATH and errors out if no compatible interpreters can be
+  found on said path.
 
   If neither variable is set, we fall back to plain PEX execution using PATH searching or the
   currently executing interpreter. If compatibility constraints are used, we match those constraints
   against these interpreters.
 
-  :param compatibility_constraints: list of requirements-style strings that constrain the
-  Python interpreter to re-exec this pex with.
+  :param compatibility_constraints: optional list of requirements-style strings that constrain the
+                                    Python interpreter to re-exec this pex with.
   """
-  if os.environ.pop('SHOULD_EXIT_BOOTSTRAP_REEXEC', None):
+
+  current_interpreter = PythonInterpreter.get()
+
+  # NB: Used only for tests.
+  if '_PEX_EXEC_CHAIN' in os.environ:
+    flag_or_chain = os.environ.pop('_PEX_EXEC_CHAIN')
+    pex_exec_chain = [] if flag_or_chain == '1' else flag_or_chain.split(os.pathsep)
+    pex_exec_chain.append(current_interpreter.binary)
+    os.environ['_PEX_EXEC_CHAIN'] = os.pathsep.join(pex_exec_chain)
+
+  current_interpreter_blessed_env_var = '_PEX_SHOULD_EXIT_BOOTSTRAP_REEXEC'
+  if os.environ.pop(current_interpreter_blessed_env_var, None):
     # We've already been here and selected an interpreter. Continue to execution.
     return
 
-  target = None
-  with TRACER.timed('Selecting runtime interpreter based on pexrc', V=3):
+  with TRACER.timed('Selecting runtime interpreter', V=3):
     if ENV.PEX_PYTHON and not ENV.PEX_PYTHON_PATH:
       # preserve PEX_PYTHON re-exec for backwards compatibility
       # TODO: Kill this off completely in favor of PEX_PYTHON_PATH
       # https://github.com/pantsbuild/pex/issues/431
+      TRACER.log('Using PEX_PYTHON={} constrained by {}'
+                 .format(ENV.PEX_PYTHON, compatibility_constraints), V=3)
       target = _select_pex_python_interpreter(ENV.PEX_PYTHON,
                                               compatibility_constraints=compatibility_constraints)
-    elif ENV.PEX_PYTHON_PATH:
-      target = _select_interpreter(pex_python_path=ENV.PEX_PYTHON_PATH,
-                                   compatibility_constraints=compatibility_constraints)
-
-    elif compatibility_constraints:
-      # Apply constraints to target using regular PATH
-      target = _select_interpreter(compatibility_constraints=compatibility_constraints)
-
-  if target and os.path.realpath(target) != os.path.realpath(sys.executable):
-    cmdline = [target] + sys.argv
-    TRACER.log('Re-executing: cmdline="%s", sys.executable="%s", PEX_PYTHON="%s", '
-               'PEX_PYTHON_PATH="%s", COMPATIBILITY_CONSTRAINTS="%s"'
-               % (cmdline, sys.executable, ENV.PEX_PYTHON, ENV.PEX_PYTHON_PATH,
-                  compatibility_constraints))
-    ENV.delete('PEX_PYTHON')
-    ENV.delete('PEX_PYTHON_PATH')
-    os.environ['SHOULD_EXIT_BOOTSTRAP_REEXEC'] = '1'
-    os.execve(target, cmdline, ENV.copy())
+    elif ENV.PEX_PYTHON_PATH or compatibility_constraints:
+      TRACER.log(
+        'Using {path} constrained by {constraints}'.format(
+          path='PEX_PYTHON_PATH={}'.format(ENV.PEX_PYTHON_PATH) if ENV.PEX_PYTHON_PATH else '$PATH',
+          constraints=compatibility_constraints
+        ),
+        V=3
+      )
+      target = _select_path_interpreter(path=ENV.PEX_PYTHON_PATH,
+                                        compatibility_constraints=compatibility_constraints)
+    else:
+      TRACER.log('Using the current interpreter {} since no constraints have been specified.'
+                 .format(sys.executable), V=3)
+      return
+
+  os.environ.pop('PEX_PYTHON', None)
+  os.environ.pop('PEX_PYTHON_PATH', None)
+
+  if target == current_interpreter:
+    TRACER.log('Using the current interpreter {} since it matches constraints.'
+               .format(sys.executable))
+    return
+
+  target_binary = target.binary
+  cmdline = [target_binary] + sys.argv
+  TRACER.log('Re-executing: cmdline="%s", sys.executable="%s", PEX_PYTHON="%s", '
+             'PEX_PYTHON_PATH="%s", COMPATIBILITY_CONSTRAINTS="%s"'
+             % (cmdline, sys.executable, ENV.PEX_PYTHON, ENV.PEX_PYTHON_PATH,
+                compatibility_constraints))
+
+  # Avoid a re-run through compatibility_constraint checking.
+  os.environ[current_interpreter_blessed_env_var] = '1'
+
+  os.execv(target_binary, cmdline)
 
 
 def _bootstrap(entry_point):

pex/variables.py

@@ -75,9 +75,6 @@ def __init__(self, environ=None, rc=None, use_defaults=True):
   def copy(self):
     return self._environ.copy()
 
-  def delete(self, variable):
-    self._environ.pop(variable, None)
-
   def set(self, variable, value):
     self._environ[variable] = str(value)
 
@@ -351,5 +348,6 @@ def PEX_EMIT_WARNINGS(self):
   def __repr__(self):
     return '{}({!r})'.format(type(self).__name__, self._environ)
 
+
 # Global singleton environment
 ENV = Variables()