Fix PythonInterpreter caching and ergonomics.

pex/interpreter.py

@@ -51,7 +51,9 @@
   )
 )
 
-setuptools_path = None
+"""
+
+EXTRAS_PY = b"""\
 try:
   import pkg_resources
 except ImportError:
@@ -69,7 +71,7 @@
 """
 
 
-def _generate_identity_source():
+def _generate_identity_source(include_site_extras):
   # Determine in the most platform-compatible way possible the identity of the interpreter
   # and its known packages.
   encodables = (
@@ -81,10 +83,11 @@ def _generate_identity_source():
     get_impl_ver
   )
 
-  return ID_PY_TMPL.replace(
-    b'__CODE__',
-    b'\n\n'.join(getsource(func).encode('utf-8') for func in encodables)
-  )
+  source = ID_PY_TMPL.replace(b'__CODE__',
+                              b'\n\n'.join(getsource(func).encode('utf-8') for func in encodables))
+  if include_site_extras:
+    source += EXTRAS_PY
+  return source
 
 
 class PythonIdentity(object):
@@ -321,25 +324,31 @@ def iter_lines():
         yield ((dist_name, dist_version), location)
     return dict(iter_lines())
 
+  @staticmethod
+  def _iter_extras(path_extras):
+    for item in path_extras:
+      for dist in find_distributions(item):
+        if dist.version:
+          yield ((dist.key, dist.version), dist.location)
+
   @classmethod
-  def _from_binary_internal(cls, path_extras):
-    def iter_extras():
-      for item in sys.path + list(path_extras):
-        for dist in find_distributions(item):
-          if dist.version:
-            yield ((dist.key, dist.version), dist.location)
-    return cls(sys.executable, PythonIdentity.get(), dict(iter_extras()))
+  def _from_binary_internal(cls, path_extras, include_site_extras):
+    extras = sys.path + list(path_extras) if include_site_extras else list(path_extras)
+    return cls(sys.executable, PythonIdentity.get(), dict(cls._iter_extras(extras)))
 
   @classmethod
-  def _from_binary_external(cls, binary, path_extras):
+  def _from_binary_external(cls, binary, path_extras, include_site_extras):
     environ = cls.sanitized_environment()
-    environ['PYTHONPATH'] = ':'.join(path_extras)
-    stdout, _ = Executor.execute([binary], env=environ, stdin_payload=_generate_identity_source())
+    stdout, _ = Executor.execute([binary],
+                                 env=environ,
+                                 stdin_payload=_generate_identity_source(include_site_extras))
     output = stdout.splitlines()
     if len(output) == 0:
       raise cls.IdentificationError('Could not establish identity of %s' % binary)
-    identity, extras = output[0], output[1:]
-    return cls(binary, PythonIdentity.from_id_string(identity), extras=cls._parse_extras(extras))
+    identity, raw_extras = output[0], output[1:]
+    extras = cls._parse_extras(raw_extras)
+    extras.update(cls._iter_extras(path_extras))
+    return cls(binary, PythonIdentity.from_id_string(identity), extras=extras)
 
   @classmethod
   def expand_path(cls, path):
@@ -366,14 +375,26 @@ def from_env(cls, hashbang):
             TRACER.log('Could not identify %s: %s' % (fn, e))
 
   @classmethod
-  def from_binary(cls, binary, path_extras=None):
+  def from_binary(cls, binary, path_extras=None, include_site_extras=True):
+    """Create an interpreter from the given `binary`.
+
+    :param str binary: The path to the python interpreter binary.
+    :param path_extras: Extra PYTHONPATH entries to add to the interpreter's `sys.path`.
+    :type path_extras: list of str
+    :param bool include_site_extras: `True` to include the `site-packages` associated
+                                     with `binary` in the interpreter's `sys.path`.
+    :return: an interpreter created from the given `binary` with only the specified
+             extras.
+    :rtype: :class:`PythonInterpreter`
+    """
     path_extras = path_extras or ()
-    if binary not in cls.CACHE:
+    key = (binary, tuple(path_extras), include_site_extras)
+    if key not in cls.CACHE:
       if binary == sys.executable:
-        cls.CACHE[binary] = cls._from_binary_internal(path_extras)
+        cls.CACHE[key] = cls._from_binary_internal(path_extras, include_site_extras)
       else:
-        cls.CACHE[binary] = cls._from_binary_external(binary, path_extras)
-    return cls.CACHE[binary]
+        cls.CACHE[key] = cls._from_binary_external(binary, path_extras, include_site_extras)
+    return cls.CACHE[key]
 
   @classmethod
   def find(cls, paths):

pex/testing.py

@@ -319,7 +319,7 @@ def bootstrap_python_installer(dest):
     raise RuntimeError("Helper method could not clone pyenv from git after 3 tries")
 
 
-def ensure_python_interpreter(version):
+def ensure_python_distribution(version):
   pyenv_root = os.path.join(os.getcwd(), '.pyenv_test')
   interpreter_location = os.path.join(pyenv_root, 'versions', version)
   pyenv = os.path.join(pyenv_root, 'bin', 'pyenv')
@@ -330,7 +330,13 @@ def ensure_python_interpreter(version):
 
   if not os.path.exists(interpreter_location):
     os.environ['PYENV_ROOT'] = pyenv_root
-    subprocess.call([pyenv, 'install', version])
-    subprocess.call([pip, 'install', SETUPTOOLS_REQUIREMENT])
+    subprocess.check_call([pyenv, 'install', '--keep', version])
+    subprocess.check_call([pip, 'install', SETUPTOOLS_REQUIREMENT])
+
+  python = os.path.join(interpreter_location, 'bin', 'python' + version[0:3])
+  return python, pip
 
-  return os.path.join(interpreter_location, 'bin', 'python' + version[0:3])
+
+def ensure_python_interpreter(version):
+  python, _ = ensure_python_distribution(version)
+  return python

tests/test_interpreter.py

@@ -2,18 +2,28 @@
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
 import os
+import subprocess
 
 import pytest
 
 from pex import interpreter
-from pex.testing import IS_PYPY, ensure_python_interpreter
+from pex.testing import (
+    IS_PYPY,
+    ensure_python_distribution,
+    ensure_python_interpreter,
+    temporary_dir
+)
 
 try:
   from mock import patch
 except ImportError:
   from unittest.mock import patch
 
 
+def version_from_tuple(version_tuple):
+  return '.'.join(str(x) for x in version_tuple)
+
+
 class TestPythonInterpreter(object):
 
   @pytest.mark.skipif('sys.version_info >= (3,0)')
@@ -23,10 +33,63 @@ def test_all_does_not_raise_with_empty_path_envvar(self):
       reload(interpreter)
       interpreter.PythonInterpreter.all()
 
+  TEST_INTERPRETER1_VERSION_TUPLE = (2, 7, 10)
+  TEST_INTERPRETER1_VERSION = version_from_tuple(TEST_INTERPRETER1_VERSION_TUPLE)
+
+  TEST_INTERPRETER2_VERSION_TUPLE = (2, 7, 9)
+  TEST_INTERPRETER2_VERSION = version_from_tuple(TEST_INTERPRETER2_VERSION_TUPLE)
+
+  @pytest.fixture
+  def test_interpreter1(self):
+    return ensure_python_interpreter(self.TEST_INTERPRETER1_VERSION)
+
+  @pytest.fixture
+  def test_interpreter2(self):
+    return ensure_python_interpreter(self.TEST_INTERPRETER2_VERSION)
+
+  @pytest.mark.skipif(IS_PYPY)
+  def test_interpreter_versioning(self, test_interpreter1):
+    py_interpreter = interpreter.PythonInterpreter.from_binary(test_interpreter1)
+    assert py_interpreter.identity.version == self.TEST_INTERPRETER1_VERSION_TUPLE
+
+  @pytest.mark.skipif(IS_PYPY)
+  def test_interpreter_caching_basic(self, test_interpreter1, test_interpreter2):
+    py_interpreter1 = interpreter.PythonInterpreter.from_binary(test_interpreter1)
+    py_interpreter2 = interpreter.PythonInterpreter.from_binary(test_interpreter2)
+    assert py_interpreter1 is not py_interpreter2
+    assert py_interpreter2.identity.version == self.TEST_INTERPRETER2_VERSION_TUPLE
+
+    py_interpreter3 = interpreter.PythonInterpreter.from_binary(test_interpreter1)
+    assert py_interpreter1 is py_interpreter3
+
+  @pytest.mark.skipif(IS_PYPY)
+  def test_interpreter_caching_include_site_extras(self, test_interpreter1):
+    py_interpreter1 = interpreter.PythonInterpreter.from_binary(test_interpreter1,
+                                                                include_site_extras=False)
+    py_interpreter2 = interpreter.PythonInterpreter.from_binary(test_interpreter1,
+                                                                include_site_extras=True)
+    py_interpreter3 = interpreter.PythonInterpreter.from_binary(test_interpreter1)
+    assert py_interpreter1 is not py_interpreter2
+    assert py_interpreter1.identity.version == py_interpreter2.identity.version
+    assert py_interpreter2 is py_interpreter3
+
   @pytest.mark.skipif(IS_PYPY)
-  def test_interpreter_versioning(self):
-    test_version_tuple = (2, 7, 10)
-    test_version = '.'.join(str(x) for x in test_version_tuple)
-    test_interpreter = ensure_python_interpreter(test_version)
-    py_interpreter = interpreter.PythonInterpreter.from_binary(test_interpreter)
-    assert py_interpreter.identity.version == test_version_tuple
+  def test_interpreter_caching_path_extras(self):
+    python, pip = ensure_python_distribution(self.TEST_INTERPRETER1_VERSION)
+    with temporary_dir() as path_extra:
+      subprocess.check_call([pip,
+                             'install',
+                             '--target={}'.format(path_extra),
+                             'ansicolors==1.1.8'])
+      py_interpreter1 = interpreter.PythonInterpreter.from_binary(python,
+                                                                  path_extras=[path_extra],
+                                                                  include_site_extras=False)
+      py_interpreter2 = interpreter.PythonInterpreter.from_binary(python,
+                                                                  include_site_extras=False)
+      py_interpreter3 = interpreter.PythonInterpreter.from_binary(python,
+                                                                  path_extras=[path_extra],
+                                                                  include_site_extras=False)
+      assert py_interpreter1 is not py_interpreter2
+      assert py_interpreter1.extras == {('ansicolors', '1.1.8'): path_extra}
+      assert py_interpreter2.extras == {}
+      assert py_interpreter1 is py_interpreter3