Add the ability to include extensions & shared libraries in pexs.

Extensions already (sort-of) worked via add_source - this makes sure that when adding a .(so|dll|pyd), we force zip_safe to False.
In addition this adds the ability to include other shared libraries with the pex. This is intended to be used for dependencies of the python extensions (e.g. a cityhash extension will depend on a cityhash library).

pex/bin/pex.py

@@ -329,6 +329,14 @@ def configure_clp():
       help='Add requirements from the given requirements file.  This option can be used multiple '
            'times.')
 
+  parser.add_option(
+      '--native-library',
+      dest='native_libraries',
+      metavar='FILE',
+      default=[],
+      action='append',
+      help='Native library to include in the archive & loader path.')
+
   parser.add_option(
       '-v',
       dest='verbosity',
@@ -477,6 +485,9 @@ def build_pex(args, options, resolver_option_builder):
     pex_builder.add_distribution(dist)
     pex_builder.add_requirement(dist.as_requirement())
 
+  for lib in options.native_libraries:
+    pex_builder.add_native_library(lib)
+
   if options.entry_point and options.script:
     die('Must specify at most one entry point or script.', INVALID_OPTIONS)
 

pex/environment.py

@@ -117,6 +117,52 @@ def __init__(self, pex, pex_info, interpreter=None, **kw):
     super(PEXEnvironment, self).__init__(
         search_path=sys.path if pex_info.inherit_path else [], **kw)
 
+  def maybe_reexec_for_library_path(self):
+    """Makes sure the files in the native-libs directory are accessible.
+
+    Because of the way native library loaders work, this function may not
+    return, and moreover, will re-launch the python interpreter.
+
+    We take advantage of the fact that pexs are deterministically extracted."""
+    if not self._pex_info.native_libraries:
+      return
+
+    if os.path.isdir(self._pex):
+      exploded_dir = self._pex
+    else:
+      exploded_dir = self.force_local(self._pex, self._pex_info)
+      assert not self._pex_info.zip_safe
+
+    assert os.path.isdir(exploded_dir)
+
+    if os.uname()[0] == 'Darwin':
+      var_name = 'DYLD_LIBRARY_PATH'
+    else:
+      var_name = 'LD_LIBRARY_PATH'
+
+    expected_path = os.environ.get(
+      'PEX_NATIVE_LIBS_PATH', os.path.join(exploded_dir, 'native-libs'))
+
+    paths = []
+    if os.environ.get(var_name):
+      paths = os.environ[var_name].split(':')
+    if expected_path in paths:
+      TRACER.log('Found native libraries in %s' % expected_path)
+      return
+
+    paths.insert(0, expected_path)
+    os.environ[var_name] = ':'.join(paths)
+
+    TRACER.log('Re-execing for %s' % var_name)
+    sys.stderr.flush()
+    sys.stdout.flush()
+    # This is technically not quite right since we may have been launched as
+    # `./foo.pex ...` and not `python ./foo.pex`. Since python doesn't
+    # seem to expose this information, there should be no harm in always
+    # using the latter form. Separately, this also means that any arguments to
+    # `python` won't be copied over.
+    os.execve(sys.executable, [sys.executable] + sys.argv, os.environ)
+
   def update_candidate_distributions(self, distribution_iter):
     for dist in distribution_iter:
       if self.can_add(dist):

pex/pex.py

@@ -56,10 +56,8 @@ def __init__(self, pex=sys.argv[0], interpreter=None, env=ENV):
     self._envs = []
     self._working_set = None
 
-  def _activate(self):
-    if not self._working_set:
-      working_set = WorkingSet([])
-
+  def _load_envs(self):
+    if not self._envs:
       # set up the local .pex environment
       pex_info = self._pex_info.copy()
       pex_info.update(self._pex_info_overrides)
@@ -71,8 +69,18 @@ def _activate(self):
         pex_info.update(self._pex_info_overrides)
         self._envs.append(PEXEnvironment(pex_path, pex_info))
 
-      # activate all of them
-      for env in self._envs:
+    return self._envs
+
+  def _maybe_reexec_for_library_path(self):
+    for env in self._load_envs():
+      env.maybe_reexec_for_library_path()
+
+  def _activate(self):
+    if not self._working_set:
+      working_set = WorkingSet([])
+
+      # activate all of the envs
+      for env in self._load_envs():
         for dist in env.activate():
           working_set.add(dist)
 
@@ -307,6 +315,9 @@ def execute(self):
     the interpreter.
     """
     try:
+      # Do this upfront so we don't waste time patching.
+      self._maybe_reexec_for_library_path()
+
       with self.patch_sys():
         working_set = self._activate()
         TRACER.log('PYTHONPATH contains:')

pex/pex_builder.py

@@ -133,11 +133,13 @@ def add_source(self, filename, env_filename):
     """
     self._ensure_unfrozen('Adding source')
     self._chroot.link(filename, env_filename, "source")
-    if filename.endswith('.py'):
+    if env_filename.endswith('.py'):
       env_filename_pyc = os.path.splitext(env_filename)[0] + '.pyc'
       with open(filename) as fp:
         pyc_object = CodeMarshaller.from_py(fp.read(), env_filename)
       self._chroot.write(pyc_object.to_pyc(), env_filename_pyc, 'source')
+    elif any(env_filename.endswith(ext) for ext in ('.so', '.dll', '.pyd')):
+      self._pex_info.add_native_library(env_filename, CacheHelper.hash(filename))
 
   def add_resource(self, filename, env_filename):
     """Add a resource to the PEX environment.
@@ -301,6 +303,25 @@ def add_egg(self, egg):
     self._ensure_unfrozen('Adding an egg')
     return self.add_dist_location(egg)
 
+  def add_native_library(self, location, name=None):
+    """Add a native library by its location on disk
+
+    :param location: The path to the native library (.so/.dylib/.dll)
+    :name: (optional) The name of the library when extracted (defaults to basename(location))
+
+    If any python extension library found in the package depends on a non-python native library,
+    this method can be used to include the native library in the pex. It will be extracted and
+    added to the proper loader path (LD_LIBRARY_PATH on most unixes).
+    """
+    self._ensure_unfrozen('Adding a native library')
+    if name is None:
+      name = os.path.basename(location)
+    target = os.path.join('native-libs', name)
+
+    self._chroot.link(location, target)
+    sha = CacheHelper.hash(location)
+    self._pex_info.add_native_library(target, sha)
+
   # TODO(wickman) Consider changing this behavior to put the onus on the consumer
   # of pex to write the pex sources correctly.
   def _prepare_inits(self):

pex/pex_info.py

@@ -28,13 +28,15 @@ class PexInfo(object):
   distributions: {dist_name: str}   # map from distribution name (i.e. path in
                                     # the internal cache) to its cache key (sha1)
   requirements: list                # list of requirements for this environment
+  native_libraries: {path: str}     # list of native libraries (extensions & dependencies)
 
   # Environment options
   pex_root: ~/.pex                   # root of all pex-related files
   entry_point: string                # entry point into this pex
   script: string                     # script to execute in this pex environment
                                      # at most one of script/entry_point can be specified
-  zip_safe: True, default False      # is this pex zip safe?
+  zip_safe: True, default True       # is this pex zip safe? (forced to False in the
+                                     # presence of native_libraries)
   inherit_path: True, default False  # should this pex inherit site-packages + PYTHONPATH?
   ignore_errors: True, default False # should we ignore inability to resolve dependencies?
   always_write_cache: False          # should we always write the internal cache to disk first?
@@ -66,6 +68,7 @@ def default(cls):
     pex_info = {
       'requirements': [],
       'distributions': {},
+      'native_libraries': {},
       'build_properties': cls.make_build_properties(),
     }
     return cls(info=pex_info)
@@ -120,6 +123,7 @@ def __init__(self, info=None):
                        '%s of type %s' % (info, type(info)))
     self._pex_info = info or {}
     self._distributions = self._pex_info.get('distributions', {})
+    self._native_libraries = self._pex_info.get('native_libraries', [])
     requirements = self._pex_info.get('requirements', [])
     if not isinstance(requirements, (list, tuple)):
       raise ValueError('Expected requirements to be a list, got %s' % type(requirements))
@@ -157,6 +161,9 @@ def zip_safe(self):
     By default zip_safe is True.  May be overridden at runtime by the $PEX_FORCE_LOCAL environment
     variable.
     """
+    if self._native_libraries:
+      return False
+
     return self._pex_info.get('zip_safe', True)
 
   @zip_safe.setter
@@ -225,6 +232,13 @@ def add_distribution(self, location, sha):
   def distributions(self):
     return self._distributions
 
+  @property
+  def native_libraries(self):
+    return self._native_libraries
+
+  def add_native_library(self, path, sha):
+    self._native_libraries[path] = sha
+
   @property
   def always_write_cache(self):
     return self._pex_info.get('always_write_cache', False)
@@ -259,6 +273,7 @@ def update(self, other):
     self._pex_info.update(other._pex_info)
     self._distributions.update(other.distributions)
     self._requirements.update(other.requirements)
+    self._native_libraries.update(other.native_libraries)
 
   def dump(self):
     pex_info_copy = self._pex_info.copy()

tests/test_integration.py

@@ -3,9 +3,11 @@
 
 import os
 
-from twitter.common.contextutil import temporary_file
+from twitter.common.contextutil import temporary_dir, temporary_file
 
-from pex.testing import run_simple_pex_test
+from pex.compatibility import nested
+from pex.pex_builder import PEXBuilder
+from pex.testing import run_simple_pex, run_simple_pex_test, temporary_content
 
 
 def test_pex_execute():
@@ -30,3 +32,30 @@ def test_pex_interpreter():
     so, rc = run_simple_pex_test("", args=(fp.name,), coverage=True, env=env)
     assert so == b'Hello world\n'
     assert rc == 0
+
+
+def test_pex_library_path():
+  with nested(temporary_dir(),
+              temporary_content({'_ext.so': 125}),
+              temporary_dir()) as (pb_dir, lib, out_dir):
+    main_py = os.path.join(pb_dir, 'exe.py')
+    with open(main_py, 'w') as pyfile:
+      pyfile.write("""
+import os
+import os.path
+
+var_name = 'DYLD_LIBRARY_PATH' if os.uname()[0] == 'Darwin' else 'LD_LIBRARY_PATH'
+print(any(os.path.exists(os.path.join(p, '_ext.so'))
+          for p in os.environ[var_name].split(':')))
+""")
+
+    pb = PEXBuilder(path=pb_dir)
+    pb.add_native_library(os.path.join(lib, '_ext.so'))
+    pb.set_executable(main_py)
+    pb.freeze()
+    assert not pb.info.zip_safe
+    pex = os.path.join(out_dir, 'app.pex')
+    pb.build(pex)
+    out, rc = run_simple_pex(pex)
+    assert rc == 0
+    assert out == b'True\n'