Extract a generalized V2 rule to inject __init__.py files (pantsbui…

…ld#7722)

### Problem
Python requires packages to have an `__init__.py` file to be recognized as a proper module for the sake of imports.

pantsbuild#7696 does this for Pytest, but inlines the logic, even though it will likely be helpful for other Python rules as well.

Further, because this logic was originally written before being able to from `Digest->Snapshot` thanks to pantsbuild#7725, we had to use `FilesContent` to grab the paths of the digest. This would mean that every single source file would be materialized and persisted to memory, resulting in extremely high memory usage (found as part of investigation in pantsbuild#7741). There is no need for the actual content, just the paths, so this is a huge inefficiency.

Will close pantsbuild#7715.

### Solution
Generalize into `@rule(InitInjectedDigest, [Snapshot])`, where `InitInjectedDigest` is a thin wrapper around a `Digest`.

We take a `Snapshot` because we need the file paths to work properly. This contrasts with earlier using `FileContents` to get the same paths. A `Snapshot` is much more light weight.

We return a `Digest` because that is the minimum information necessary to work properly, and the caller of the rule can then convert that `Digest` back into a `Snapshot`.

### Result
It will now be easier for other Python rules to work with Python packages.

The unnecessary memory usage is now fixed. The V2 Pytest runner now has a space complexity of `O(t + t*e + b)`, rather than `O(t + t*e + s)`, where `t` is # targets, `e` is # env vars, `b` is # `BUILD` files, and `s` is # source files.

src/python/pants/backend/python/register.py

@@ -8,7 +8,7 @@
 from pants.backend.python.python_artifact import PythonArtifact
 from pants.backend.python.python_requirement import PythonRequirement
 from pants.backend.python.python_requirements import PythonRequirements
-from pants.backend.python.rules import python_test_runner
+from pants.backend.python.rules import inject_init, python_test_runner
 from pants.backend.python.targets.python_app import PythonApp
 from pants.backend.python.targets.python_binary import PythonBinary
 from pants.backend.python.targets.python_distribution import PythonDistribution
@@ -81,4 +81,4 @@ def register_goals():
 
 
 def rules():
-  return tuple(python_test_runner.rules())
+  return inject_init.rules() + python_test_runner.rules()

src/python/pants/backend/python/rules/BUILD

@@ -8,5 +8,6 @@ python_library(
     'src/python/pants/engine:selectors',
     'src/python/pants/rules/core',
     'src/python/pants/source:source',
+    'src/python/pants/util:objects',
   ],
 )

src/python/pants/backend/python/rules/inject_init.py

@@ -0,0 +1,44 @@
+# coding=utf-8
+# Copyright 2019 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+from pants.backend.python.subsystems.pex_build_util import identify_missing_init_files
+from pants.engine.fs import EMPTY_DIRECTORY_DIGEST, Digest, Snapshot
+from pants.engine.isolated_process import ExecuteProcessRequest, ExecuteProcessResult
+from pants.engine.rules import rule
+from pants.engine.selectors import Get
+from pants.util.objects import datatype
+
+
+# TODO(#7710): Once this gets fixed, rename this to InitInjectedDigest.
+class InjectedInitDigest(datatype([('directory_digest', Digest)])): pass
+
+
+@rule(InjectedInitDigest, [Snapshot])
+def inject_init(snapshot):
+  """Ensure that every package has an __init__.py file in it."""
+  missing_init_files = tuple(sorted(identify_missing_init_files(snapshot.files)))
+  if not missing_init_files:
+    new_init_files_digest = EMPTY_DIRECTORY_DIGEST
+  else:
+    # TODO(7718): add a builtin rule for FilesContent->Snapshot, so that we can avoid using touch
+    # and the absolute path and have the engine build the files for us.
+    touch_init_request = ExecuteProcessRequest(
+      argv=("/usr/bin/touch",) + missing_init_files,
+      output_files=missing_init_files,
+      description="Inject missing __init__.py files: {}".format(", ".join(missing_init_files)),
+      input_files=snapshot.directory_digest,
+    )
+    touch_init_result = yield Get(ExecuteProcessResult, ExecuteProcessRequest, touch_init_request)
+    new_init_files_digest = touch_init_result.output_directory_digest
+  # TODO(#7710): Once this gets fixed, merge the original source digest and the new init digest
+  # into one unified digest.
+  yield InjectedInitDigest(directory_digest=new_init_files_digest)
+
+
+def rules():
+  return [
+      inject_init,
+    ]

src/python/pants/backend/python/rules/python_test_runner.py

@@ -10,11 +10,11 @@
 
 from future.utils import text_type
 
-from pants.backend.python.subsystems.pex_build_util import identify_missing_init_files
+from pants.backend.python.rules.inject_init import InjectedInitDigest
 from pants.backend.python.subsystems.pytest import PyTest
 from pants.backend.python.subsystems.python_setup import PythonSetup
-from pants.engine.fs import (Digest, DirectoriesToMerge, DirectoryWithPrefixToStrip, FilesContent,
-                             Snapshot, UrlToFetch)
+from pants.engine.fs import (Digest, DirectoriesToMerge, DirectoryWithPrefixToStrip, Snapshot,
+                             UrlToFetch)
 from pants.engine.isolated_process import (ExecuteProcessRequest, ExecuteProcessResult,
                                            FallibleExecuteProcessResult)
 from pants.engine.legacy.graph import BuildFileAddresses, TransitiveHydratedTargets
@@ -128,25 +128,13 @@ def run_python_test(test_target, pytest, python_setup, source_root_config):
     Digest, DirectoriesToMerge(directories=tuple(all_sources_digests)),
   )
 
-  # TODO(7716): add a builtin rule to go from DirectoriesToMerge->Snapshot or Digest->Snapshot.
-  # TODO(7715): generalize the injection of __init__.py files.
-  # TODO(7718): add a builtin rule for FilesContent->Snapshot.
-  file_contents = yield Get(FilesContent, Digest, sources_digest)
-  file_paths = [fc.path for fc in file_contents]
-  injected_inits = tuple(sorted(identify_missing_init_files(file_paths)))
-  if injected_inits:
-    touch_init_request = ExecuteProcessRequest(
-      argv=("/usr/bin/touch",) + injected_inits,
-      output_files=injected_inits,
-      description="Inject empty __init__.py into all packages without one already.",
-      input_files=sources_digest,
-    )
-    touch_init_result = yield Get(ExecuteProcessResult, ExecuteProcessRequest, touch_init_request)
-
-  all_input_digests = [sources_digest, requirements_pex_response.output_directory_digest]
-  if injected_inits:
-    all_input_digests.append(touch_init_result.output_directory_digest)
+  inits_digest = yield Get(InjectedInitDigest, Digest, sources_digest)
 
+  all_input_digests = [
+    sources_digest,
+    inits_digest.directory_digest,
+    requirements_pex_response.output_directory_digest,
+  ]
   merged_input_files = yield Get(
     Digest,
     DirectoriesToMerge,

tests/python/pants_test/backend/python/rules/BUILD

@@ -1,6 +1,18 @@
 # Copyright 2019 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
+python_tests(
+  name='inject_init',
+  sources=['test_inject_init.py'],
+  dependencies=[
+    'src/python/pants/backend/python/rules',
+    'src/python/pants/engine:fs',
+    'src/python/pants/engine:rules',
+    'src/python/pants/util:collections',
+    'tests/python/pants_test:test_base',
+  ]
+)
+
 python_tests(
   name='python_test_runner',
   sources=['test_python_test_runner.py'],

tests/python/pants_test/backend/python/rules/test_inject_init.py

@@ -0,0 +1,39 @@
+# coding=utf-8
+# Copyright 2019 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+from pants.backend.python.rules.inject_init import InjectedInitDigest, inject_init
+from pants.engine.fs import EMPTY_DIRECTORY_DIGEST, EMPTY_SNAPSHOT, Snapshot
+from pants.engine.rules import RootRule
+from pants.util.collections import assert_single_element
+from pants_test.test_base import TestBase
+
+
+class TestInjectInit(TestBase):
+
+  @classmethod
+  def rules(cls):
+    return super(TestInjectInit, cls).rules() + [inject_init, RootRule(Snapshot)]
+
+  def assert_result(self, input_snapshot, expected_digest):
+    injected_digest = assert_single_element(
+      self.scheduler.product_request(InjectedInitDigest, [input_snapshot])
+    )
+    self.assertEqual(injected_digest.directory_digest, expected_digest)
+
+  def test_noops_when_empty_snapshot(self):
+    self.assert_result(input_snapshot=EMPTY_SNAPSHOT, expected_digest=EMPTY_DIRECTORY_DIGEST)
+
+  def test_noops_when_init_already_present(self):
+    snapshot = self.make_snapshot({
+      "test/foo.py": "",
+      "test/__init__.py": ""
+    })
+    self.assert_result(input_snapshot=snapshot, expected_digest=EMPTY_DIRECTORY_DIGEST)
+
+  def test_adds_when_init_missing(self):
+    snapshot = self.make_snapshot({"test/foo.py": ""})
+    expected_digest = self.make_snapshot({"test/__init__.py": ""}).directory_digest
+    self.assert_result(input_snapshot=snapshot, expected_digest=expected_digest)