Improve PEX_PATH handling in pex.py (#421)

Replace all usages of self._vars.PEX_PATH with a unified pex path composed from merging pex_path in PEX-INFO and PEX_PATH set in the environment.

Problem:
pex_info.pex_path was being overwritten in _activate() in pex.py by the path set in the environment.

Solution:
Implement a util function to merge the paths into a single pex path for usage throughout pex.py, giving priority to the PEX-INFO pex path if both are set.

pex/pex.py

@@ -21,7 +21,7 @@
 from .orderedset import OrderedSet
 from .pex_info import PexInfo
 from .tracer import TRACER
-from .util import iter_pth_paths
+from .util import iter_pth_paths, merge_split
 from .variables import ENV
 
 
@@ -69,10 +69,14 @@ def _activate(self):
       # set up the local .pex environment
       pex_info = self._pex_info.copy()
       pex_info.update(self._pex_info_overrides)
+      pex_info.merge_pex_path(self._vars.PEX_PATH)
       self._envs.append(PEXEnvironment(self._pex, pex_info))
-
+      # N.B. by this point, `pex_info.pex_path` will contain a single pex path
+      # merged from pex_path in `PEX-INFO` and `PEX_PATH` set in the environment.
+      # `PEX_PATH` entries written into `PEX-INFO` take precedence over those set
+      # in the environment.
       if pex_info.pex_path:
-        # set up other environments as specified in PEX_PATH
+        # set up other environments as specified in pex_path
         for pex_path in filter(None, pex_info.pex_path.split(os.pathsep)):
           pex_info = PexInfo.from_pex(pex_path)
           pex_info.update(self._pex_info_overrides)
@@ -279,7 +283,7 @@ def patch_all(path, path_importer_cache, modules):
         sys.path[:], sys.path_importer_cache.copy(), sys.modules.copy())
     new_sys_path, new_sys_path_importer_cache, new_sys_modules = self.minimum_sys(inherit_path)
 
-    new_sys_path.extend(filter(None, self._vars.PEX_PATH.split(os.pathsep)))
+    new_sys_path.extend(merge_split(self._pex_info.pex_path, self._vars.PEX_PATH))
 
     patch_all(new_sys_path, new_sys_path_importer_cache, new_sys_modules)
     yield

pex/pex_info.py

@@ -12,6 +12,7 @@
 from .compatibility import string as compatibility_string
 from .compatibility import PY2
 from .orderedset import OrderedSet
+from .util import merge_split
 from .variables import ENV
 
 PexPlatform = namedtuple('PexPlatform', 'interpreter version strict')
@@ -98,7 +99,6 @@ def from_env(cls, env=ENV):
       'inherit_path': supplied_env.PEX_INHERIT_PATH,
       'ignore_errors': supplied_env.PEX_IGNORE_ERRORS,
       'always_write_cache': supplied_env.PEX_ALWAYS_CACHE,
-      'pex_path': supplied_env.PEX_PATH,
     }
     # Filter out empty entries not explicitly set in the environment.
     return cls(info=dict((k, v) for (k, v) in pex_info.items() if v is not None))
@@ -284,3 +284,11 @@ def dump(self, **kwargs):
 
   def copy(self):
     return self.from_json(self.dump())
+
+  def merge_pex_path(self, pex_path):
+    """Merges a new PEX_PATH definition into the existing one (if any).
+    :param string pex_path: The PEX_PATH to merge.
+    """
+    if not pex_path:
+      return
+    self.pex_path = ':'.join(merge_split(self.pex_path, pex_path))

pex/util.py

@@ -245,3 +245,14 @@ def iter_pth_paths(filename):
         if extras_dir_case_insensitive not in known_paths and os.path.exists(extras_dir):
           yield extras_dir
           known_paths.add(extras_dir_case_insensitive)
+
+
+def merge_split(*paths):
+  """Merge paths into a single path delimited by colons and split on colons to return
+  a list of paths.
+
+  :param paths: a variable length list of path strings
+  :return: a list of paths from the merged path list split by colons
+  """
+  filtered_paths = filter(None, paths)
+  return [p for p in ':'.join(filtered_paths).split(':') if p]

tests/test_integration.py

@@ -278,3 +278,51 @@ def test_pex_path_arg():
     stdout, rc = run_simple_pex(pex_out_path, [test_file_path])
     assert rc == 0
     assert stdout == b'Success!\n'
+
+
+def test_pex_path_in_pex_info_and_env():
+  with temporary_dir() as output_dir:
+
+    # create 2 pex files for PEX-INFO pex_path
+    pex1_path = os.path.join(output_dir, 'pex1.pex')
+    res1 = run_pex_command(['--disable-cache', 'requests', '-o', pex1_path])
+    res1.assert_success()
+    pex2_path = os.path.join(output_dir, 'pex2.pex')
+    res2 = run_pex_command(['--disable-cache', 'flask', '-o', pex2_path])
+    res2.assert_success()
+    pex_path = ':'.join(os.path.join(output_dir, name) for name in ('pex1.pex', 'pex2.pex'))
+
+    # create a pex for environment PEX_PATH
+    pex3_path = os.path.join(output_dir, 'pex3.pex')
+    res3 = run_pex_command(['--disable-cache', 'wheel', '-o', pex3_path])
+    res3.assert_success()
+    env_pex_path = os.path.join(output_dir, 'pex3.pex')
+
+    # parameterize the pex arg for test.py
+    pex_out_path = os.path.join(output_dir, 'out.pex')
+    # create test file test.py that attempts to import modules from pex1/pex2
+    test_file_path = os.path.join(output_dir, 'test.py')
+    with open(test_file_path, 'w') as fh:
+      fh.write(dedent('''
+        import requests
+        import flask
+        import wheel
+        import sys
+        import os
+        import subprocess
+        print('Success!')
+        '''))
+
+    # build out.pex composed from pex1/pex1
+    run_pex_command(['--disable-cache',
+      '--pex-path={}'.format(pex_path),
+      '-o', pex_out_path])
+
+    # load secondary PEX_PATH
+    env = os.environ.copy()
+    env['PEX_PATH'] = env_pex_path
+
+    # run test.py with composite env
+    stdout, rc = run_simple_pex(pex_out_path, [test_file_path], env=env)
+    assert rc == 0
+    assert stdout == b'Success!\n'

tests/test_util.py

@@ -13,7 +13,13 @@
 from pex.installer import EggInstaller, WheelInstaller
 from pex.pex_builder import PEXBuilder
 from pex.testing import make_bdist, run_simple_pex, temporary_content, write_zipfile
-from pex.util import CacheHelper, DistributionHelper, iter_pth_paths, named_temporary_file
+from pex.util import (
+    CacheHelper,
+    DistributionHelper,
+    iter_pth_paths,
+    merge_split,
+    named_temporary_file
+)
 
 try:
   from unittest import mock
@@ -204,3 +210,17 @@ def test_iter_pth_paths(mock_exists):
       with open(pth_tmp_path, 'wb') as f:
         f.write(to_bytes(pth_content))
       assert sorted(PTH_TEST_MAPPING[pth_content]) == sorted(list(iter_pth_paths(pth_tmp_path)))
+
+
+def test_merge_split():
+  path_1, path_2 = '/pex/path/1:/pex/path/2', '/pex/path/3:/pex/path/4'
+  result = merge_split(path_1, path_2)
+  assert result == ['/pex/path/1', '/pex/path/2', '/pex/path/3', '/pex/path/4']
+
+  path_1, path_2 = '/pex/path/1:', '/pex/path/3:/pex/path/4'
+  result = merge_split(path_1, path_2)
+  assert result == ['/pex/path/1', '/pex/path/3', '/pex/path/4']
+
+  path_1, path_2 = '/pex/path/1::/pex/path/2', '/pex/path/3:/pex/path/4'
+  result = merge_split(path_1, path_2)
+  assert result == ['/pex/path/1', '/pex/path/2', '/pex/path/3', '/pex/path/4']