jsmith/access zipped assets

pex/common.py

@@ -279,7 +279,10 @@ def write(self, data, dst, label=None, mode='wb'):
     self._tag(dst, label)
     self._mkdir_for(dst)
     with open(os.path.join(self.chroot, dst), mode) as wp:
-      wp.write(data)
+      try:
+        wp.write(data)
+      except TypeError:
+        wp.write(bytes(data, 'UTF-8'))
 
   def touch(self, dst, label=None):
     """Perform 'touch' on {chroot}/dest with optional label.

pex/environment.py

@@ -13,10 +13,13 @@
     Environment,
     find_distributions,
     Requirement,
+    resource_isdir,
+    resource_listdir,
+    resource_string,
     WorkingSet
 )
 
-from .common import open_zip, safe_mkdir, safe_rmtree
+from .common import open_zip, safe_mkdir, safe_mkdtemp, safe_rmtree
 from .interpreter import PythonInterpreter
 from .package import distribution_compatible
 from .pex_builder import PEXBuilder
@@ -98,6 +101,33 @@ def load_internal_cache(cls, pex, pex_info):
         for dist in cls.write_zipped_internal_cache(pex, pex_info):
           yield dist
 
+  @classmethod
+  def access_zipped_assets(cls, static_module_name, static_path, asset_path, dir_location='.'):
+    """
+    Create a copy of static resource files as we can't serve
+    them from within the pex file.
+
+    :param static_module_name: Module name containing module to cache in a tempdir
+    :type static_module_name: string in the form of 'twitter.common.zookeeper'
+    :param static_path: Module name of the form 'serverset'
+    :param asset_path: Initially a module name that's the same as the static_path, but will be changed to walk
+        the directory tree
+    :param dir_location: directory to create a new temporary directory in
+    """
+    temp_dir = safe_mkdtemp(dir=dir_location)
+    for asset in resource_listdir(static_module_name, asset_path):
+      asset_target = os.path.join(os.path.relpath(asset_path, static_path), asset)[2:]
+      if resource_isdir(static_module_name, os.path.join(asset_path, asset)):
+        safe_mkdir(os.path.join(temp_dir, asset_target))
+        cls.access_zipped_assets(static_module_name, static_path, os.path.join(asset_path, asset))
+      else:
+        with open(os.path.join(temp_dir, asset_target), 'wb') as fp:
+          path = os.path.join(static_path, asset_target)
+          file_data = resource_string(static_module_name, path)
+          fp.write(file_data)
+    return temp_dir
+
+
   def __init__(self, pex, pex_info, interpreter=None, **kw):
     self._internal_cache = os.path.join(pex, pex_info.internal_cache)
     self._pex = pex

pex/testing.py

@@ -122,6 +122,14 @@ def make_bdist(name='my_project', installer_impl=EggInstaller, zipped=False, zip
 
 
 def write_simple_pex(td, exe_contents, dists=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 coverage: include coverage header
+  """
   dists = dists or []
 
   with open(os.path.join(td, 'exe.py'), 'w') as fp:

tests/test_environment.py

@@ -3,14 +3,22 @@
 
 import os
 from contextlib import contextmanager
-
+import subprocess
+from textwrap import dedent
+
+try:
+  import mock
+except ImportError:
+  from unittest import mock
+import pkg_resources
 from twitter.common.contextutil import temporary_dir, temporary_file
 
+from pex.common import safe_mkdir, safe_mkdtemp
 from pex.compatibility import nested
 from pex.environment import PEXEnvironment
 from pex.pex_builder import PEXBuilder
 from pex.pex_info import PexInfo
-from pex.testing import make_bdist
+from pex.testing import make_bdist, run_simple_pex_test
 
 
 @contextmanager
@@ -93,3 +101,66 @@ def test_load_internal_cache_unzipped():
     assert len(dists) == 1
     assert normalize(dists[0].location).startswith(
         normalize(os.path.join(pb.path(), pb.info.internal_cache)))
+
+@mock.patch('pex.environment.resource_string', spec=pkg_resources.resource_string)
+@mock.patch('pex.environment.resource_isdir', spec=pkg_resources.resource_isdir)
+@mock.patch('pex.environment.resource_listdir', spec=pkg_resources.resource_listdir)
+def test_access_zipped_assets(mock_resource_listdir, mock_resource_isdir, mock_resource_string):
+  try:
+    import __builtin__
+    builtin_path = '__builtin__'
+  except ImportError:
+    # Python3
+    import builtins
+    builtin_path = 'builtins'
+
+  mock_open = mock.mock_open()
+  with mock.patch('%s.open' % builtin_path, mock_open, create=True):
+    mock_resource_listdir.side_effect = [['./__init__.py', './directory/'], ['file.py']]
+    mock_resource_isdir.side_effect = [False, True, False]
+    mock_resource_string.return_value = 'testing'
+
+    PEXEnvironment.access_zipped_assets('twitter.common', 'dirutil', 'dirutil')
+
+    assert mock_resource_listdir.call_count == 2
+    assert mock_open.call_count == 2
+    file_handle = mock_open.return_value.__enter__.return_value
+    assert file_handle.write.call_count == 2
+
+def test_access_zipped_assets_integration():
+  test_executable = dedent('''
+      import os
+      from _pex.environment import PEXEnvironment
+      temp_dir = PEXEnvironment.access_zipped_assets('my_package', 'submodule', 'submodule')
+      with open(os.path.join(temp_dir, 'mod.py'), 'r') as fp:
+        for line in fp:
+          print(line)
+  ''')
+  with nested(temporary_dir(), temporary_dir()) as (td1, td2):
+    pb = PEXBuilder(path=td1)
+    with open(os.path.join(td1, 'exe.py'), 'w') as fp:
+      fp.write(test_executable)
+      pb.set_executable(fp.name)
+
+    submodule = os.path.join(td1, 'my_package', 'submodule')
+    safe_mkdir(submodule)
+    mod_path = os.path.join(submodule, 'mod.py')
+    with open(mod_path, 'w') as fp:
+      fp.write('accessed')
+      pb.add_source(fp.name, 'my_package/submodule/mod.py')
+
+    pex = os.path.join(td2, 'app.pex')
+    pb.build(pex)
+
+    po = subprocess.Popen(
+        [pex],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT)
+    po.wait()
+    output = po.stdout.read()
+    try:
+      output = output.decode('UTF-8')
+    except:
+      pass
+    assert output == 'accessed\n'
+    assert po.returncode == 0