Repository: pex-tool/pex

Issue:
Pex should not lose PEX_PYTHON or PEX_PYTHON_PATH when re-exec-ing
Currently Pex loses PEX_PYTHON and PEX_PYTHON_PATH constraints on re-exec. The bug stems from #441 as commented there (https://github.com/pantsbuild/pex/pull/441/files#r277150236) in the following code:
https://github.com/pantsbuild/pex/blob/28869f6af835a39eae96453c0e5ab3dc01f7b00e/pex/pex_bootstrapper.py#L89-L136

In particular note lines 135 and 136. We re-exec vie `execve`, passing an explicit environment. That environment dict is obtained from `ENV.copy()` which does not see the `os.environ` mutation on the prior line:
https://github.com/pantsbuild/pex/blob/ffd56ce35cec240469567e2a51710d62baa21f08/pex/variables.py#L67-L76



Discussion:
Noticed working pantsbuild/pants#7591
I can jump in on this. My approach is to replace line 135 with usage of https://github.com/pantsbuild/pex/blob/ffd56ce35cec240469567e2a51710d62baa21f08/pex/variables.py#L81 to set that variable, and then yank out the PP and PPP deletion lines. Thoughts?
Sorry, missed the in progress label.
> ... to set that variable, and then yank out the PP and PPP deletion lines. Thoughts?

`Variables` really ought to be immutable (or as much as possible anyway). I took the opposite tack and just mutated the current process environment before re-exec'ing:
```python
os.environ.pop('PEX_PYTHON', None)
os.environ.pop('PEX_PYTHON_PATH', None)
os.environ['SHOULD_EXIT_BOOTSTRAP_REEXEC'] = '1'

os.execv(target, cmdline)
```

Resolving patch:
diff --git a/pex/interpreter.py b/pex/interpreter.py
--- a/pex/interpreter.py
+++ b/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:
diff --git a/pex/pex_bootstrapper.py b/pex/pex_bootstrapper.py
--- a/pex/pex_bootstrapper.py
+++ b/pex/pex_bootstrapper.py
@@ -18,38 +18,44 @@
 __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
@@ -57,83 +63,114 @@ def find_compatible_interpreters(pex_python_path=None, compatibility_constraints
   )
 
 
-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):
diff --git a/pex/variables.py b/pex/variables.py
--- a/pex/variables.py
+++ b/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()


Test patch:
diff --git a/tests/test_integration.py b/tests/test_integration.py
--- a/tests/test_integration.py
+++ b/tests/test_integration.py
@@ -3,6 +3,7 @@
 
 import filecmp
 import functools
+import json
 import os
 import platform
 import subprocess
@@ -43,7 +44,10 @@
 
 def make_env(**kwargs):
   env = os.environ.copy()
-  env.update((k, str(v)) for k, v in kwargs.items())
+  env.update((k, str(v)) for k, v in kwargs.items() if v is not None)
+  for k, v in kwargs.items():
+    if v is None:
+      env.pop(k, None)
   return env
 
 
@@ -603,7 +607,8 @@ def test_pex_python():
     stdin_payload = b'import sys; print(sys.executable); sys.exit(0)'
     stdout, rc = run_simple_pex(pex_out_path, stdin=stdin_payload, env=env)
     assert rc == 1
-    fail_str = 'not compatible with specified interpreter constraints'.encode()
+    fail_str = ('Failed to find a compatible PEX_PYTHON={} for constraints'
+                .format(pex_python)).encode()
     assert fail_str in stdout
 
     # test PEX_PYTHON with no constraints
@@ -1482,3 +1487,84 @@ def test_issues_736_requirement_setup_py_with_extras():
         env=make_env(PEX_INTERPRETER='1')
       )
       assert output.decode('utf-8').strip() == u'hello world!'
+
+
+def _assert_exec_chain(exec_chain=None,
+                       pex_python=None,
+                       pex_python_path=None,
+                       interpreter_constraints=None):
+
+  with temporary_dir() as td:
+    test_pex = os.path.join(td, 'test.pex')
+
+    args = ['-o', test_pex]
+    if interpreter_constraints:
+      args.extend('--interpreter-constraint={}'.format(ic) for ic in interpreter_constraints)
+
+    env = os.environ.copy()
+    PATH = env.get('PATH').split(os.pathsep)
+
+    def add_to_path(entry):
+      if os.path.isfile(entry):
+        entry = os.path.dirname(entry)
+      PATH.append(entry)
+
+    if pex_python:
+      add_to_path(pex_python)
+    if pex_python_path:
+      for path in pex_python_path:
+        add_to_path(path)
+
+    env['PATH'] = os.pathsep.join(PATH)
+    result = run_pex_command(args, env=env)
+    result.assert_success()
+
+    env = make_env(_PEX_EXEC_CHAIN=1,
+                   PEX_INTERPRETER=1,
+                   PEX_PYTHON=pex_python,
+                   PEX_PYTHON_PATH=os.pathsep.join(pex_python_path) if pex_python_path else None)
+
+    output = subprocess.check_output(
+      [sys.executable, test_pex, '-c', 'import json, os; print(json.dumps(os.environ.copy()))'],
+      env=env
+    )
+    final_env = json.loads(output.decode('utf-8'))
+
+    assert 'PEX_PYTHON' not in final_env
+    assert 'PEX_PYTHON_PATH' not in final_env
+    assert '_PEX_SHOULD_EXIT_BOOTSTRAP_REEXEC' not in final_env
+
+    expected_exec_chain = [os.path.realpath(i) for i in [sys.executable] + (exec_chain or [])]
+    assert expected_exec_chain == final_env['_PEX_EXEC_CHAIN'].split(os.pathsep)
+
+
+def test_pex_no_reexec_no_constraints():
+  _assert_exec_chain()
+
+
+def test_pex_no_reexec_constraints_match_current():
+  current_version = '.'.join(str(component) for component in sys.version_info[0:3])
+  _assert_exec_chain(interpreter_constraints=['=={}'.format(current_version)])
+
+
+def test_pex_reexec_constraints_dont_match_current_pex_python_path():
+  py36_interpreter = ensure_python_interpreter(PY36)
+  py27_interpreter = ensure_python_interpreter(PY27)
+  _assert_exec_chain(exec_chain=[py36_interpreter],
+                     pex_python_path=[py27_interpreter, py36_interpreter],
+                     interpreter_constraints=['=={}'.format(PY36)])
+
+
+def test_pex_reexec_constraints_dont_match_current_pex_python_path_min_py_version_selected():
+  py36_interpreter = ensure_python_interpreter(PY36)
+  py27_interpreter = ensure_python_interpreter(PY27)
+  _assert_exec_chain(exec_chain=[py27_interpreter],
+                     pex_python_path=[py36_interpreter, py27_interpreter])
+
+
+def test_pex_reexec_constraints_dont_match_current_pex_python():
+  version = PY27 if sys.version_info[0:2] == (3, 6) else PY36
+  interpreter = ensure_python_interpreter(version)
+  _assert_exec_chain(exec_chain=[interpreter],
+                     pex_python=interpreter,
+                     interpreter_constraints=['=={}'.format(version)])
diff --git a/tests/test_pex_bootstrapper.py b/tests/test_pex_bootstrapper.py
--- a/tests/test_pex_bootstrapper.py
+++ b/tests/test_pex_bootstrapper.py
@@ -17,7 +17,7 @@ def test_find_compatible_interpreters():
 
   def find_interpreters(*constraints):
     return [interp.binary for interp in
-            find_compatible_interpreters(pex_python_path=pex_python_path,
+            find_compatible_interpreters(path=pex_python_path,
                                          compatibility_constraints=constraints)]
 
   assert [py35, py36] == find_interpreters('>3')
diff --git a/tests/test_variables.py b/tests/test_variables.py
--- a/tests/test_variables.py
+++ b/tests/test_variables.py
@@ -110,11 +110,10 @@ def test_pex_vars_hermetic():
 
 def test_pex_vars_set():
   v = Variables(environ={})
+  assert v._get_int('HELLO') is None
   v.set('HELLO', '42')
   assert v._get_int('HELLO') == 42
-  v.delete('HELLO')
-  assert v._get_int('HELLO') is None
-  assert {} == v.copy()
+  assert {'HELLO': '42'} == v.copy()
 
 
 def test_pex_get_kv():