From 96e27a0a86d81482a8a2f0c996e5d14ad7592ac3 Mon Sep 17 00:00:00 2001
From: Benjy Weinberger <benjyw@gmail.com>
Date: Wed, 7 Dec 2016 19:15:39 -0800
Subject: [PATCH] Make PEX_PATH unify pex sources, as well as requirements.

---
 pex/pex.py        |  7 ++++---
 pex/testing.py    | 13 ++++++++++++-
 tests/test_pex.py | 32 ++++++++++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/pex/pex.py b/pex/pex.py
index 138984fad..753806c65 100644
--- a/pex/pex.py
+++ b/pex/pex.py
@@ -233,9 +233,8 @@ def patch(working_set):
   # potentially in a wonky state since the patches here (minimum_sys_modules
   # for example) actually mutate global state.  This should not be
   # considered a reversible operation despite being a contextmanager.
-  @classmethod
   @contextmanager
-  def patch_sys(cls, inherit_path):
+  def patch_sys(self, inherit_path):
     """Patch sys with all site scrubbed."""
     def patch_dict(old_value, new_value):
       old_value.clear()
@@ -248,7 +247,9 @@ def patch_all(path, path_importer_cache, modules):
 
     old_sys_path, old_sys_path_importer_cache, old_sys_modules = (
         sys.path[:], sys.path_importer_cache.copy(), sys.modules.copy())
-    new_sys_path, new_sys_path_importer_cache, new_sys_modules = cls.minimum_sys(inherit_path)
+    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)))
 
     patch_all(new_sys_path, new_sys_path_importer_cache, new_sys_modules)
     yield
diff --git a/pex/testing.py b/pex/testing.py
index 5f993228b..fad0a777b 100644
--- a/pex/testing.py
+++ b/pex/testing.py
@@ -150,16 +150,20 @@ def make_bdist(name='my_project', installer_impl=EggInstaller, zipped=False, zip
 """
 
 
-def write_simple_pex(td, exe_contents, dists=None, coverage=False):
+def write_simple_pex(td, exe_contents, dists=None, sources=None, coverage=False):
   """Write a pex file that contains an executable entry point
 
   :param td: temporary directory path
   :param exe_contents: entry point python file
   :type exe_contents: string
   :param dists: distributions to include, typically sdists or bdists
+  :param sources: sources to include, as a list of pairs (env_filename, contents)
   :param coverage: include coverage header
   """
   dists = dists or []
+  sources = sources or []
+
+  safe_mkdir(td)
 
   with open(os.path.join(td, 'exe.py'), 'w') as fp:
     fp.write(exe_contents)
@@ -169,6 +173,13 @@ def write_simple_pex(td, exe_contents, dists=None, coverage=False):
   for dist in dists:
     pb.add_egg(dist.location)
 
+  for env_filename, contents in sources:
+    src_path = os.path.join(td, env_filename)
+    safe_mkdir(os.path.dirname(src_path))
+    with open(src_path, 'w') as fp:
+      fp.write(contents)
+    pb.add_source(src_path, env_filename)
+
   pb.set_executable(os.path.join(td, 'exe.py'))
   pb.freeze()
 
diff --git a/tests/test_pex.py b/tests/test_pex.py
index 2560445ac..356f76afb 100644
--- a/tests/test_pex.py
+++ b/tests/test_pex.py
@@ -206,3 +206,35 @@ def test_pex_run():
 
       fake_stdout.seek(0)
       assert fake_stdout.read() == b'hellohello'
+
+
+def test_pex_paths():
+  # Tests that PEX_PATH allows importing sources from the referenced pex.
+  with named_temporary_file() as fake_stdout:
+    with temporary_dir() as temp_dir:
+      pex1_path = os.path.join(temp_dir, 'pex1')
+      write_simple_pex(
+        pex1_path,
+        exe_contents='',
+        sources=[
+          ('foo_pkg/__init__.py', ''),
+          ('foo_pkg/foo_module.py', 'def foo_func():\n  return "42"')
+        ]
+      )
+
+      pex2_path = os.path.join(temp_dir, 'pex2')
+      pex2 = write_simple_pex(
+        pex2_path,
+        'import sys; from bar_pkg.bar_module import bar_func; '
+        'sys.stdout.write(bar_func()); sys.exit(0)',
+        sources=[
+          ('bar_pkg/bar_module.py',
+           'from foo_pkg.foo_module import foo_func\ndef bar_func():\n  return foo_func()')
+        ]
+      )
+
+      rc = PEX(pex2.path()).run(stdin=None, stdout=fake_stdout, env={'PEX_PATH': pex1_path})
+      assert rc == 0
+
+      fake_stdout.seek(0)
+      assert fake_stdout.read() == b'42'