Generate a single python source chroot.

This change reverts the majority of #5400 (GatherSources) and adjusts
the new coverage plugin path mapping system to handle mapping a single
source chroot back to potentially multiple source roots. The existing
tests of this  capability introduced in #5400 still pass and as a result
the previously conflicting use cases represented by #5314 (multiple
library/binary source roots with interdependencies under test) and
#5426 (parallel test packages in a test source root) are now both
satisfied.

Fixes #5314
Fixes #5426
Depends on #5534

src/python/pants/backend/python/__init__.py

@@ -0,0 +1 @@
+__import__('pkg_resources').declare_namespace(__name__)

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

@@ -0,0 +1 @@
+__import__('pkg_resources').declare_namespace(__name__)

src/python/pants/backend/python/tasks/coverage/plugin.py

@@ -53,42 +53,43 @@ def no_branch_lines(self):
 class MyPlugin(CoveragePlugin):
   """A plugin that knows how to map Pants PEX chroots back to repo source code when reporting."""
 
-  def __init__(self, buildroot, src_to_chroot):
+  def __init__(self, buildroot, src_chroot_path, src_to_target_base):
     super(MyPlugin, self).__init__()
     self._buildroot = buildroot
-    self._src_to_chroot = src_to_chroot
+    self._src_chroot_path = src_chroot_path
+    self._src_to_target_base = src_to_target_base
+    self._target_bases = set(self._src_to_target_base.values())
 
   def find_executable_files(self, top):
-    for chroot_path in self._src_to_chroot.values():
-      if top.startswith(chroot_path):
-        for root, _, files in os.walk(top):
-          for f in files:
-            if f.endswith('.py'):
-              yield os.path.join(root, f)
-        break
+    if top.startswith(self._src_chroot_path):
+      for root, dirs, files in os.walk(top):
+        for f in files:
+          if f.endswith('.py'):
+            yield os.path.join(root, f)
 
   def file_tracer(self, filename):
-    for chroot_path in self._src_to_chroot.values():
-      if filename.startswith(chroot_path):
+    if filename.startswith(self._src_chroot_path):
+      src = os.path.relpath(filename, self._src_chroot_path)
+      if src in self._src_to_target_base:
         return MyFileTracer(filename)
 
   def file_reporter(self, filename):
-    src_file = self._map_to_src(filename)
-    mapped_relpath = os.path.relpath(src_file, self._buildroot)
-    return MyFileReporter(filename, mapped_relpath)
+    mapped_relpath = self._map_relpath(filename)
+    return MyFileReporter(filename, mapped_relpath or filename)
 
-  def _map_to_src(self, chroot):
-    for src_dir, chroot_dir in self._src_to_chroot.items():
-      if chroot.startswith(chroot_dir):
-        return src_dir + chroot[len(chroot_dir):]
-    raise AssertionError('Failed to map traced file {} to any source root via '
-                         'source root -> chroot mappings:\n\t{}'
-                         .format(chroot, '\n\t'.join(sorted('{} -> {}'.format(src_dir, chroot_dir)
-                                                            for src_dir, chroot_dir
-                                                            in self._src_to_chroot.items()))))
+  def _map_relpath(self, filename):
+    src = os.path.relpath(filename, self._src_chroot_path)
+    target_base = self._src_to_target_base.get(src) or self._find_target_base(src)
+    return os.path.join(target_base, src) if target_base else filename
+
+  def _find_target_base(self, src):
+    for target_base in self._target_bases:
+      if os.path.isfile(os.path.join(self._buildroot, target_base, src)):
+        return target_base
 
 
 def coverage_init(reg, options):
   buildroot = options['buildroot']
-  src_to_chroot = json.loads(options['src_to_chroot'])
-  reg.add_file_tracer(MyPlugin(buildroot, src_to_chroot))
+  src_chroot_path = options['src_chroot_path']
+  src_to_target_base = json.loads(options['src_to_target_base'])
+  reg.add_file_tracer(MyPlugin(buildroot, src_chroot_path, src_to_target_base))

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

@@ -11,6 +11,7 @@
 from pex.interpreter import PythonInterpreter
 from pex.pex import PEX
 from pex.pex_builder import PEXBuilder
+from twitter.common.collections import OrderedSet
 
 from pants.backend.python.tasks.pex_build_util import (dump_sources, has_python_sources,
                                                        has_resources, is_python_target)
@@ -22,44 +23,20 @@
 class GatherSources(Task):
   """Gather local Python sources.
 
-  Creates one or more (unzipped) PEXs on disk containing the local Python sources.
-  These PEXes can be merged with a requirements PEX to create a unified Python environment
-  for running the relevant python code.
+  Creates an (unzipped) PEX on disk containing the local Python sources. This PEX can be merged
+  with a requirements PEX to create a unified Python environment for running the relevant python
+  code.
   """
 
-  class PythonSources(object):
-    """A mapping of unzipped source PEXs by the targets whose sources the PEXs contain."""
-
-    class UnmappedTargetError(Exception):
-      """Indicates that no python source pex could be found for a given target."""
-
-    def __init__(self, pex_by_target_base):
-      self._pex_by_target_base = pex_by_target_base
-
-    def for_target(self, target):
-      """Return the unzipped PEX containing the given target's sources.
-
-      :returns: An unzipped PEX containing at least the given target's sources.
-      :rtype: :class:`pex.pex.PEX`
-      :raises: :class:`GatherSources.PythonSources.UnmappedTargetError` if no pex containing the
-               given target's sources could be found.
-      """
-      pex = self._pex_by_target_base.get(target.target_base)
-      if pex is None:
-        raise self.UnmappedTargetError()
-      return pex
-
-    def all(self):
-      """Return all the unzipped source PEXs needed for this round."""
-      return self._pex_by_target_base.values()
+  PYTHON_SOURCES = 'python_sources'
 
   @classmethod
   def implementation_version(cls):
-    return super(GatherSources, cls).implementation_version() + [('GatherSources', 4)]
+    return super(GatherSources, cls).implementation_version() + [('GatherSources', 5)]
 
   @classmethod
   def product_types(cls):
-    return [cls.PythonSources]
+    return [cls.PYTHON_SOURCES]
 
   @classmethod
   def prepare(cls, options, round_manager):
@@ -68,54 +45,25 @@ def prepare(cls, options, round_manager):
 
   def execute(self):
     interpreter = self.context.products.get_data(PythonInterpreter)
+    targets = self._collect_source_targets()
 
-    pex_by_target_base = OrderedDict()  # Preserve ~PYTHONPATH ordering over pexes.
-    for target_base, targets in self._iter_targets_by_base():
-      with self.invalidated(targets) as invalidation_check:
-        pex = self._get_pex_for_versioned_targets(interpreter, invalidation_check.all_vts)
-        pex_by_target_base[target_base] = pex
-    self.context.products.register_data(self.PythonSources, self.PythonSources(pex_by_target_base))
-
-  def _iter_targets_by_base(self):
-    # N.B: Files and Resources targets belong with the consuming (dependee) targets so that those
-    # targets can be ensured of access to the files in their PEX chroot. This means a given Files
-    # or Resources target could be embedded in multiple pexes.
+    with self.invalidated(targets) as invalidation_check:
+      pex = self._get_pex_for_versioned_targets(interpreter, invalidation_check.all_vts)
+      self.context.products.register_data(self.PYTHON_SOURCES, pex)
 
-    context = self.context
-    python_target_addresses = [p.address for p in context.targets(predicate=is_python_target)]
+  def _collect_source_targets(self):
+    python_target_addresses = [p.address for p in self.context.targets(predicate=is_python_target)]
 
-    targets_by_base = OrderedDict()  # Preserve ~PYTHONPATH ordering over source roots.
-    resource_targets = set()
+    targets = OrderedSet()
 
     def collect_source_targets(target):
-      if has_python_sources(target):
-        targets = targets_by_base.get(target.target_base)
-        if targets is None:
-          targets = set()
-          targets_by_base[target.target_base] = targets
+      if has_python_sources(target) or has_resources(target):
         targets.add(target)
-      elif has_resources(target):
-        resource_targets.add(target)
-
-    build_graph = context.build_graph
-    build_graph.walk_transitive_dependency_graph(addresses=python_target_addresses,
-                                                 work=collect_source_targets)
-
-    for resource_target in resource_targets:
-      dependees = build_graph.transitive_dependees_of_addresses([resource_target.address])
-      for target_base, targets in targets_by_base.items():
-        for dependee in dependees:
-          if dependee in targets:
-            # N.B.: This can add the resource to too many pexes. A canonical example is
-            # test -> lib -> resource where test and lib have separate source roots. In this case
-            # the resource is added to both the test pex and the lib pex and it's only needed in the
-            # lib pex. The upshot is we allow python code access to undeclared (ie: indirect)
-            # resource dependencies which is no worse than historical precedent, but could be
-            # improved with a more complex algorithm.
-            targets.add(resource_target)
-            break
-
-    return targets_by_base.items()
+
+    self.context.build_graph.walk_transitive_dependency_graph(addresses=python_target_addresses,
+                                                              work=collect_source_targets)
+
+    return targets
 
   def _get_pex_for_versioned_targets(self, interpreter, versioned_targets):
     if versioned_targets:

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

@@ -18,6 +18,7 @@
 
 from six import StringIO
 from six.moves import configparser
+from twitter.common.collections import OrderedSet
 
 from pants.backend.python.targets.python_tests import PythonTests
 from pants.backend.python.tasks.gather_sources import GatherSources
@@ -183,7 +184,7 @@ def _ensure_section(cp, section):
 
   # N.B.: Extracted for tests.
   @classmethod
-  def _add_plugin_config(cls, cp, src_to_chroot):
+  def _add_plugin_config(cls, cp, src_chroot_path, src_to_target_base):
     # We use a coverage plugin to map PEX chroot source paths back to their original repo paths for
     # report output.
     plugin_module = 'pants.backend.python.tasks.coverage.plugin'
@@ -192,16 +193,14 @@ def _add_plugin_config(cls, cp, src_to_chroot):
 
     cp.add_section(plugin_module)
     cp.set(plugin_module, 'buildroot', get_buildroot())
-    cp.set(plugin_module,
-           'src_to_chroot',
-           json.dumps({os.path.join(get_buildroot(), f): os.path.join(get_buildroot(), t)
-                       for f, t in src_to_chroot.items()}))
+    cp.set(plugin_module, 'src_chroot_path', src_chroot_path)
+    cp.set(plugin_module, 'src_to_target_base', json.dumps(src_to_target_base))
 
-  def _generate_coverage_config(self, src_to_chroot):
+  def _generate_coverage_config(self, src_to_target_base):
     cp = configparser.SafeConfigParser()
     cp.readfp(StringIO(self.DEFAULT_COVERAGE_CONFIG))
 
-    self._add_plugin_config(cp, src_to_chroot)
+    self._add_plugin_config(cp, self._source_chroot_path, src_to_target_base)
 
     # See the debug options here: http://nedbatchelder.com/code/coverage/cmd.html#cmd-run-debug
     if self._debug:
@@ -216,8 +215,8 @@ def _generate_coverage_config(self, src_to_chroot):
     return cp
 
   @contextmanager
-  def _cov_setup(self, workdirs, coverage_morfs, src_to_chroot):
-    cp = self._generate_coverage_config(src_to_chroot=src_to_chroot)
+  def _cov_setup(self, workdirs, coverage_morfs, src_to_target_base):
+    cp = self._generate_coverage_config(src_to_target_base=src_to_target_base)
     # Note that it's important to put the tmpfile under the workdir, because pytest
     # uses all arguments that look like paths to compute its rootdir, and we want
     # it to pick the buildroot.
@@ -239,15 +238,15 @@ def _maybe_emit_coverage_data(self, workdirs, test_targets, pex):
       yield []
       return
 
-    def pex_src_root(tgt):
-      return os.path.relpath(self._source_chroot_path((tgt,)), get_buildroot())
+    pex_src_root = os.path.relpath(self._source_chroot_path, get_buildroot())
 
-    src_to_chroot = {}
+    src_to_target_base = {}
     for target in test_targets:
       libs = (tgt for tgt in target.closure()
               if tgt.has_sources('.py') and not isinstance(tgt, PythonTests))
       for lib in libs:
-        src_to_chroot[lib.target_base] = pex_src_root(lib)
+        for src in lib.sources_relative_to_source_root():
+          src_to_target_base[src] = lib.target_base
 
     def ensure_trailing_sep(path):
       return path if path.endswith(os.path.sep) else path + os.path.sep
@@ -287,13 +286,13 @@ def packages():
           rel_source = ensure_trailing_sep(rel_source)
 
           found_target_base = False
-          for target_base, pex_root in src_to_chroot.items():
+          for target_base in set(src_to_target_base.values()):
             prefix = ensure_trailing_sep(target_base)
             if rel_source.startswith(prefix):
               # ... rel_source will match on prefix=src/python/ ...
               suffix = rel_source[len(prefix):]
               # ... suffix will equal foo/bar ...
-              coverage_morfs.append(os.path.join(get_buildroot(), pex_root, suffix))
+              coverage_morfs.append(os.path.join(get_buildroot(), pex_src_root, suffix))
               found_target_base = True
               # ... and we end up appending <pex_src_root>/foo/bar to the coverage_sources.
               break
@@ -305,7 +304,7 @@ def packages():
 
     with self._cov_setup(workdirs,
                          coverage_morfs=coverage_morfs,
-                         src_to_chroot=src_to_chroot) as (args, coverage_rc):
+                         src_to_target_base=src_to_target_base) as (args, coverage_rc):
       try:
         yield args
       finally:
@@ -320,7 +319,7 @@ def coverage_run(subcommand, arguments):
 
         # The '.coverage' data file is output in the CWD of the test run above; so we make sure to
         # look for it there.
-        with self._maybe_run_in_chroot(test_targets):
+        with self._maybe_run_in_chroot():
           # On failures or timeouts, the .coverage file won't be written.
           if not os.path.exists('.coverage'):
             self.context.log.warn('No .coverage file was found! Skipping coverage reporting.')
@@ -506,7 +505,7 @@ def _do_run_tests_with_args(self, pex, args):
       return PytestResult.exception()
 
   def _map_relsrc_to_targets(self, targets):
-    pex_src_root = os.path.relpath(self._source_chroot_path(targets), get_buildroot())
+    pex_src_root = os.path.relpath(self._source_chroot_path, get_buildroot())
     # First map chrooted sources back to their targets.
     relsrc_to_target = {os.path.join(pex_src_root, src): target for target in targets
       for src in target.sources_relative_to_source_root()}
@@ -557,17 +556,8 @@ def iter_partitions():
         for test_target in test_targets:
           yield (test_target,)
     else:
-      targets_by_target_base = OrderedDict()
-      for test_target in test_targets:
-        targets_for_base = targets_by_target_base.get(test_target.target_base)
-        if targets_for_base is None:
-          targets_for_base = []
-          targets_by_target_base[test_target.target_base] = targets_for_base
-        targets_for_base.append(test_target)
-
       def iter_partitions():
-        for test_targets in targets_by_target_base.values():
-          yield tuple(test_targets)
+        yield tuple(test_targets)
 
     workdir = self.workdir
 
@@ -633,13 +623,11 @@ def _run_pytest(self, fail_fast, test_targets, workdirs):
     if not test_targets:
       return PytestResult.rc(0)
 
-    test_chroot_path = self._source_chroot_path(test_targets)
-
     # Absolute path to chrooted test file -> Path to original test file relative to the buildroot.
     sources_map = OrderedDict()
     for t in test_targets:
       for p in t.sources_relative_to_source_root():
-        sources_map[os.path.join(test_chroot_path, p)] = os.path.join(t.target_base, p)
+        sources_map[os.path.join(self._source_chroot_path, p)] = os.path.join(t.target_base, p)
 
     if not sources_map:
       return PytestResult.rc(0)
@@ -678,7 +666,7 @@ def _run_pytest(self, fail_fast, test_targets, workdirs):
       if os.path.exists(junitxml_path):
         os.unlink(junitxml_path)
 
-      with self._maybe_run_in_chroot(test_targets):
+      with self._maybe_run_in_chroot():
         result = self._do_run_tests_with_args(pytest_binary.pex, args)
 
       # There was a problem prior to test execution preventing junit xml file creation so just let
@@ -704,17 +692,9 @@ def parse_error_handler(parse_error):
 
       return result.with_failed_targets(failed_targets)
 
-  @memoized_method
-  def _source_chroot_path(self, targets):
-    if len(targets) > 1:
-      target_bases = {target.target_base for target in targets}
-      assert len(target_bases) == 1, ('Expected targets to live in the same source root, given '
-                                      'targets living under the following source roots: {}'
-                                      .format(', '.join(sorted(target_bases))))
-    representative_target = targets[0]
-
-    python_sources = self.context.products.get_data(GatherSources.PythonSources)
-    return python_sources.for_target(representative_target).path()
+  @memoized_property
+  def _source_chroot_path(self):
+    return self.context.products.get_data(GatherSources.PYTHON_SOURCES).path()
 
   def _pex_run(self, pex, workunit_name, args, env):
     with self.context.new_workunit(name=workunit_name,
@@ -724,9 +704,9 @@ def _pex_run(self, pex, workunit_name, args, env):
       return process.wait()
 
   @contextmanager
-  def _maybe_run_in_chroot(self, targets):
+  def _maybe_run_in_chroot(self):
     if self.run_tests_in_chroot:
-      with pushd(self._source_chroot_path(targets)):
+      with pushd(self._source_chroot_path):
         yield
     else:
       yield