Fix PEX_* env stripping and allow turning off.

Previously, `PEX_*` environment variable stripping incorrectly stripped
`PEX_*` environment variables in host process when the `PEX.run` API was
used. This is fixed and an option to not strip pex environment variables
at all is added to `PexInfo` / the Pex CLI. This new option is used to
publish a Pex PEX that supports control via environment variables.

Fixes pex-tool#180
Fixes pex-tool#925
Work towards pex-tool#926

pex/bin/pex.py

@@ -302,8 +302,16 @@ def configure_clp_pex_options(parser):
       '--runtime-pex-root',
       dest='runtime_pex_root',
       default=None,
-      help='Specify the pex root to be used in the generated .pex file. [Default: ~/.pex]'
-  )
+      help='Specify the pex root to be used in the generated .pex file. [Default: ~/.pex]')
+
+  group.add_option(
+      '--strip-pex-env', '--no-strip-pex-env',
+      dest='strip_pex_env',
+      default=True,
+      action='callback',
+      callback=parse_bool,
+      help='Strip all `PEX_*` environment variables used to control the PEX runtime before handing '
+           'off control to the PEX entrypoint. [Default: %default]')
 
   parser.add_option_group(group)
 
@@ -576,6 +584,7 @@ def walk_and_do(fn, src_dir):
   pex_info.emit_warnings = options.emit_warnings
   pex_info.inherit_path = options.inherit_path
   pex_info.pex_root = options.runtime_pex_root
+  pex_info.strip_pex_env = options.strip_pex_env
   if options.interpreter_constraint:
     for ic in options.interpreter_constraint:
       pex_builder.add_interpreter_constraint(ic)

pex/pex.py

@@ -43,15 +43,18 @@ class NotFound(Error): pass
   class InvalidEntryPoint(Error): pass
 
   @classmethod
-  def clean_environment(cls):
+  def _clean_environment(cls, env=None, strip_pex_env=True):
+    env = env or os.environ
+
     try:
-      del os.environ['MACOSX_DEPLOYMENT_TARGET']
+      del env['MACOSX_DEPLOYMENT_TARGET']
     except KeyError:
       pass
-    # Cannot change dictionary size during __iter__
-    filter_keys = [key for key in os.environ if key.startswith('PEX_')]
-    for key in filter_keys:
-      del os.environ[key]
+
+    if strip_pex_env:
+      for key in list(env):
+        if key.startswith('PEX_'):
+          del env[key]
 
   def __init__(self, pex=sys.argv[0], interpreter=None, env=ENV, verify_entry_point=False):
     self._pex = pex
@@ -416,7 +419,7 @@ def execute(self):
   def _execute(self):
     force_interpreter = self._vars.PEX_INTERPRETER
 
-    self.clean_environment()
+    self._clean_environment(strip_pex_env=self._pex_info.strip_pex_env)
 
     if force_interpreter:
       TRACER.log('PEX_INTERPRETER specified, dropping into interpreter')
@@ -560,7 +563,7 @@ def cmdline(self, args=()):
     cmds.extend(args)
     return cmds
 
-  def run(self, args=(), with_chroot=False, blocking=True, setsid=False, **kwargs):
+  def run(self, args=(), with_chroot=False, blocking=True, setsid=False, env=None, **kwargs):
     """Run the PythonEnvironment in an interpreter in a subprocess.
 
     :keyword args: Additional arguments to be passed to the application being invoked by the
@@ -569,10 +572,16 @@ def run(self, args=(), with_chroot=False, blocking=True, setsid=False, **kwargs)
     :keyword blocking: If true, return the return code of the subprocess.
       If false, return the Popen object of the invoked subprocess.
     :keyword setsid: If true, run the PEX in a separate operating system session.
-
+    :keyword env: An optional environment dict to use as the PEX subprocess environment. If none is
+                  passed, the ambient environment is inherited.
     Remaining keyword arguments are passed directly to subprocess.Popen.
     """
-    self.clean_environment()
+    if env is not None:
+      # If explicit env vars are passed, we don't want clean any of these.
+      env = env.copy()
+    else:
+      env = os.environ.copy()
+      self._clean_environment(env=env)
 
     cmdline = self.cmdline(args)
     TRACER.log('PEX.run invoking %s' % ' '.join(cmdline))
@@ -582,6 +591,7 @@ def run(self, args=(), with_chroot=False, blocking=True, setsid=False, **kwargs)
                                     stdin=kwargs.pop('stdin', None),
                                     stdout=kwargs.pop('stdout', None),
                                     stderr=kwargs.pop('stderr', None),
+                                    env=env,
                                     **kwargs)
     return process.wait() if blocking else process
 

pex/pex_info.py

@@ -171,6 +171,19 @@ def zip_safe(self):
   def zip_safe(self, value):
     self._pex_info['zip_safe'] = bool(value)
 
+  @property
+  def strip_pex_env(self):
+    """Whether or not this PEX should strip `PEX_*` env vars before executing its entrypoint.
+
+    You might want to set this to `False` if this PEX executes other PEXes or the Pex CLI itself and
+    you want the executed PEX to be controlled via PEX environment variables.
+    """
+    return self._pex_info.get('strip_pex_env', True)
+
+  @strip_pex_env.setter
+  def strip_pex_env(self, value):
+    self._pex_info['strip_pex_env'] = bool(value)
+
   @property
   def pex_path(self):
     """A colon separated list of other pex files to merge into the runtime environment.

pex/testing.py

@@ -10,10 +10,10 @@
 import subprocess
 import sys
 from collections import namedtuple
+from contextlib import contextmanager
 from textwrap import dedent
 
 from pex.common import open_zip, safe_mkdir, safe_mkdtemp, safe_rmtree, temporary_dir, touch
-from pex.compatibility import PY3, nested
 from pex.distribution_target import DistributionTarget
 from pex.executor import Executor
 from pex.interpreter import PythonInterpreter
@@ -307,12 +307,11 @@ def run_simple_pex(pex, args=(), interpreter=None, stdin=None, **kwargs):
                   stderr=subprocess.STDOUT,
                   **kwargs)
   stdout, _ = process.communicate(input=stdin)
-  print(stdout.decode('utf-8') if PY3 else stdout)
   return stdout.replace(b'\r', b''), process.returncode
 
 
 def run_simple_pex_test(body, args=(), env=None, dists=None, coverage=False, interpreter=None):
-  with nested(temporary_dir(), temporary_dir()) as (td1, td2):
+  with temporary_dir() as td1, temporary_dir() as td2:
     pb = write_simple_pex(td1, body, dists=dists, coverage=coverage, interpreter=interpreter)
     pex = os.path.join(td2, 'app.pex')
     pb.build(pex)
@@ -380,3 +379,21 @@ def ensure_python_distribution(version):
 def ensure_python_interpreter(version):
   python, _ = ensure_python_distribution(version)
   return python
+
+
+@contextmanager
+def environment_as(**kwargs):
+  existing = {key: os.environ.get(key) for key in kwargs}
+
+  def adjust_environment(mapping):
+    for key, value in mapping.items():
+      if value is not None:
+        os.environ[key] = value
+      else:
+        del os.environ[key]
+
+  adjust_environment(kwargs)
+  try:
+    yield
+  finally:
+    adjust_environment(existing)

scripts/package.py

@@ -34,6 +34,7 @@ def build_pex_pex() -> None:
     '--no-use-system-time',
     '--interpreter-constraint', python_requires(),
     '--python-shebang', '/usr/bin/env python',
+    '--no-strip-pex-env',
     '-o', 'dist/pex',
     '-c', 'pex',
     pex_requirement

tests/test_integration.py

@@ -1685,3 +1685,44 @@ def test_issues_898():
       assert zipp_location.startswith(pex_root), (
         'Failed to import zipp from {} under {}'.format(pex_file, python)
       )
+
+
+def test_pex_run_strip_env(tmpdir):
+  src_dir = os.path.join(tmpdir, 'src')
+  with safe_open(os.path.join(src_dir, 'print_pex_env.py'), 'w') as fp:
+    fp.write(dedent("""
+      import json
+      import os
+
+      print(json.dumps({k: v for k, v in os.environ.items() if k.startswith('PEX_')}))
+    """))
+
+  pex_env = dict(PEX_ROOT=os.path.join(tmpdir, 'pex_root'))
+  env = make_env(**pex_env)
+
+  stripped_pex_file = os.path.join(tmpdir, 'stripped.pex')
+  results = run_pex_command(
+    args=[
+      '--sources-directory={}'.format(src_dir),
+      '--entry-point=print_pex_env',
+      '-o', stripped_pex_file
+    ],
+  )
+  results.assert_success()
+  assert {} == json.loads(subprocess.check_output([stripped_pex_file], env=env)), (
+    'Expected the entrypoint environment to be stripped of PEX_ environment variables.'
+  )
+
+  unstripped_pex_file = os.path.join(tmpdir, 'unstripped.pex')
+  results = run_pex_command(
+    args=[
+      '--sources-directory={}'.format(src_dir),
+      '--entry-point=print_pex_env',
+      '--no-strip-pex-env',
+      '-o', unstripped_pex_file
+    ],
+  )
+  results.assert_success()
+  assert pex_env == json.loads(subprocess.check_output([unstripped_pex_file], env=env)), (
+    'Expected the entrypoint environment to be left un-stripped.'
+  )

tests/test_pex.py

@@ -1,13 +1,15 @@
 # Copyright 2014 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
+import json
 import os
 import subprocess
 import sys
 import tempfile
 import textwrap
 from collections import namedtuple
 from contextlib import contextmanager
+from textwrap import dedent
 from types import ModuleType
 
 import pytest
@@ -25,6 +27,7 @@
     WheelBuilder,
     built_wheel,
     ensure_python_interpreter,
+    environment_as,
     make_bdist,
     named_temporary_file,
     run_simple_pex,
@@ -570,6 +573,32 @@ def test_execute_interpreter_file_program():
       assert b'' == stderr
 
 
+def test_pex_run_strip_env():
+  with temporary_dir() as pex_root:
+    pex_env = dict(PEX_MODULE='does_not_exist_in_sub_pex', PEX_ROOT=pex_root)
+    with environment_as(**pex_env), temporary_dir() as pex_chroot:
+      pex_builder = PEXBuilder(path=pex_chroot)
+      with tempfile.NamedTemporaryFile() as fp:
+        fp.write(dedent(b"""
+          import json
+          import os
+
+          print(json.dumps({k: v for k, v in os.environ.items() if k.startswith("PEX_")}))
+        """))
+        fp.flush()
+        pex_builder.set_executable(fp.name, 'print_pex_env.py')
+      pex_builder.freeze()
+
+      stdout, returncode = run_simple_pex(pex_chroot)
+      assert 0 == returncode
+      assert {} == json.loads(stdout), (
+        'Expected the entrypoint environment to be stripped of PEX_ environment variables.'
+      )
+      assert pex_env == {k: v for k, v in os.environ.items() if k.startswith("PEX_")}, (
+        'Expected the parent environment to be left un-stripped.'
+      )
+
+
 def test_pex_run_custom_setuptools_useable():
   resolved_dists = resolve(['setuptools==36.2.7'])
   dists = [resolved_dist.distribution for resolved_dist in resolved_dists]

tests/test_variables.py

@@ -1,10 +1,11 @@
 # Copyright 2015 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
+
 import os
-from contextlib import contextmanager
 
 import pytest
 
+from pex.testing import environment_as
 from pex.util import named_temporary_file
 from pex.variables import Variables
 
@@ -69,24 +70,6 @@ def test_pex_get_int():
     assert Variables(environ={'HELLO': 'welp'})._get_int('HELLO')
 
 
-@contextmanager
-def environment_as(**kwargs):
-  existing = {key: os.environ.get(key) for key in kwargs}
-
-  def adjust_environment(mapping):
-    for key, value in mapping.items():
-      if value is not None:
-        os.environ[key] = value
-      else:
-        del os.environ[key]
-
-  adjust_environment(kwargs)
-  try:
-    yield
-  finally:
-    adjust_environment(existing)
-
-
 def assert_pex_vars_hermetic():
   v = Variables()
   assert os.environ == v.copy()