Move construction of PythonChroot to PythonTask base class.

Previously each subtask was creating its own.
This commit replaces temporary_pex_builder with temporary_chroot.

Isolating PythonChroot creation in just one place will make it easier
to transition to python-environment-as-a-subsystem.

Testing Done:
CI passes: https://travis-ci.org/pantsbuild/pants/builds/55242907.

Reviewed at https://rbcommons.com/s/twitter/r/1965/

src/python/pants/backend/python/tasks/python_binary_create.py

@@ -8,7 +8,6 @@
 import os
 import time
 
-from pants.backend.python.python_chroot import PythonChroot
 from pants.backend.python.targets.python_binary import PythonBinary
 from pants.backend.python.tasks.python_task import PythonTask
 from pants.base.exceptions import TaskError
@@ -49,14 +48,6 @@ def create_binary(self, binary):
     pexinfo = binary.pexinfo.copy()
     pexinfo.build_properties = build_properties
 
-    with self.temporary_pex_builder(pex_info=pexinfo, interpreter=interpreter) as builder:
-      chroot = PythonChroot(
-        context=self.context,
-        targets=[binary],
-        builder=builder,
-        platforms=binary.platforms,
-        interpreter=interpreter)
-
+    with self.temporary_chroot(interpreter=interpreter, pex_info=pexinfo, targets=[binary], platforms=binary.platforms) as chroot:
       pex_path = os.path.join(self._distdir, '%s.pex' % binary.name)
-      chroot.dump()
-      builder.build(pex_path)
+      chroot.builder.build(pex_path)

src/python/pants/backend/python/tasks/python_eval.py

@@ -10,7 +10,6 @@
 
 from pex.pex import PEX
 
-from pants.backend.python.python_chroot import PythonChroot
 from pants.backend.python.targets.python_binary import PythonBinary
 from pants.backend.python.targets.python_library import PythonLibrary
 from pants.backend.python.tasks.python_task import PythonTask
@@ -134,29 +133,19 @@ def _compile_target(self, target):
       else:
         pexinfo, platforms = None, None
 
-      with self.temporary_pex_builder(interpreter=interpreter, pex_info=pexinfo) as builder:
-        with self.context.new_workunit(name='resolve'):
-          chroot = PythonChroot(
-              context=self.context,
-              targets=[target],
-              builder=builder,
-              platforms=platforms,
-              interpreter=interpreter)
-
-          chroot.dump()
-
-        with temporary_file() as imports_file:
+      with temporary_file() as imports_file:
+        def pre_freeze(chroot):
           generator = Generator(pkgutil.get_data(__name__, self._EVAL_TEMPLATE_PATH),
                                 chroot=chroot.path(),
                                 modules=modules)
           generator.write(imports_file)
           imports_file.close()
+          chroot.builder.set_executable(imports_file.name, '__pants_python_eval__.py')
 
-          builder.set_executable(imports_file.name, '__pants_python_eval__.py')
-
-          builder.freeze()
-          pex = PEX(builder.path(), interpreter=interpreter)
-
+        with self.temporary_chroot(interpreter=interpreter, pex_info=pexinfo,
+                                   targets=[target], platforms=platforms,
+                                   pre_freeze=pre_freeze) as chroot:
+          pex = PEX(chroot.builder.path(), interpreter=interpreter)
           with self.context.new_workunit(name='eval',
                                          labels=[WorkUnit.COMPILER, WorkUnit.RUN, WorkUnit.TOOL],
                                          cmd=' '.join(pex.cmdline())) as workunit:

src/python/pants/backend/python/tasks/python_repl.py

@@ -7,7 +7,6 @@
 
 from pex.pex import PEX
 
-from pants.backend.python.python_chroot import PythonChroot
 from pants.backend.python.python_requirement import PythonRequirement
 from pants.backend.python.tasks.python_task import PythonTask
 from pants.base.target import Target
@@ -45,18 +44,10 @@ def execute(self, **pex_run_kwargs):
       else:
         entry_point = 'code:interact'
 
-      with self.temporary_pex_builder(interpreter=interpreter) as builder:
-        builder.set_entry_point(entry_point)
-        chroot = PythonChroot(
-          context=self.context,
-          targets=targets,
-          extra_requirements=extra_requirements,
-          builder=builder,
-          interpreter=interpreter)
-
-        chroot.dump()
-        builder.freeze()
-        pex = PEX(builder.path(), interpreter=interpreter)
+      with self.temporary_chroot(interpreter=interpreter, targets=targets,
+                                 extra_requirements=extra_requirements,
+                                 pre_freeze=lambda ch: ch.builder.set_entry_point(entry_point)) as chroot:
+        pex = PEX(chroot.builder.path(), interpreter=interpreter)
         self.context.release_lock()
         with stty_utils.preserve_stty_settings():
           with self.context.new_workunit(name='run', labels=[WorkUnit.RUN]):

src/python/pants/backend/python/tasks/python_run.py

@@ -9,7 +9,6 @@
 
 from pex.pex import PEX
 
-from pants.backend.python.python_chroot import PythonChroot
 from pants.backend.python.targets.python_binary import PythonBinary
 from pants.backend.python.tasks.python_task import PythonTask
 from pants.base.exceptions import TaskError
@@ -37,17 +36,8 @@ def execute(self):
       # jvm_binary, in which case we have to no-op and let jvm_run do its thing.
       # TODO(benjy): Some more elegant way to coordinate how tasks claim targets.
       interpreter = self.select_interpreter_for_targets(self.context.targets())
-      with self.temporary_pex_builder(interpreter=interpreter, pex_info=binary.pexinfo) as builder:
-        chroot = PythonChroot(
-          context=self.context,
-          targets=[binary],
-          builder=builder,
-          platforms=binary.platforms,
-          interpreter=interpreter)
-
-        chroot.dump()
-        builder.freeze()
-        pex = PEX(builder.path(), interpreter=interpreter)
+      with self.temporary_chroot(interpreter=interpreter, pex_info=binary.pexinfo, targets=[binary], platforms=binary.platforms) as chroot:
+        pex = PEX(chroot.builder.path(), interpreter=interpreter)
         self.context.release_lock()
         with self.context.new_workunit(name='run', labels=[WorkUnit.RUN]):
           args = []

src/python/pants/backend/python/tasks/python_task.py

@@ -13,6 +13,7 @@
 
 from pants.backend.core.tasks.task import Task
 from pants.backend.python.interpreter_cache import PythonInterpreterCache
+from pants.backend.python.python_chroot import PythonChroot
 from pants.backend.python.python_setup import PythonRepos, PythonSetup
 from pants.base.exceptions import TaskError
 
@@ -83,9 +84,26 @@ def select_interpreter(self, filters):
     return interpreter
 
   @contextmanager
-  def temporary_pex_builder(self, interpreter=None, pex_info=None, parent_dir=None):
-    """Yields a PEXBuilder and cleans up its chroot when it goes out of context."""
-    path = tempfile.mkdtemp(dir=parent_dir)
+  def temporary_chroot(self, interpreter=None, pex_info=None, targets=None,
+                       extra_requirements=None, platforms=None, pre_freeze=None):
+    """Yields a temporary PythonChroot created with the specified args.
+
+    pre_freeze is an optional function run on the chroot just before freezing its builder,
+    to allow for any extra modification.
+    """
+    path = tempfile.mkdtemp()
     builder = PEXBuilder(path=path, interpreter=interpreter, pex_info=pex_info)
-    yield builder
-    builder.chroot().delete()
+    with self.context.new_workunit('chroot'):
+      chroot = PythonChroot(
+        context=self.context,
+        targets=targets,
+        extra_requirements=extra_requirements,
+        builder=builder,
+        platforms=platforms,
+        interpreter=interpreter)
+      chroot.dump()
+      if pre_freeze:
+        pre_freeze(chroot)
+      builder.freeze()
+    yield chroot
+    chroot.delete()