Inline inject_init.py into prepare_chrooted_python_sources.py (#9115

)

We were only using `inject_init.py` for the purpose of `prepare_chrooted_python_sources.py`. Inlining this simplifies things. If we ever end up needing the init logic to be used in another context, we can, of course, generalize it back out.

More importantly, this stops using a hacky `ExecuteProcessRequest` to create the missing `__init__.py` files to instead use engine intrinsics. This improves readability and performance.

src/python/pants/backend/awslambda/python/awslambda_python_rules_test.py

@@ -10,7 +10,6 @@
 from pants.backend.awslambda.python.awslambda_python_rules import rules as awslambda_python_rules
 from pants.backend.python.rules import (
   download_pex_bin,
-  inject_init,
   pex,
   pex_from_target_closure,
   prepare_chrooted_python_sources,
@@ -42,7 +41,6 @@ def rules(cls):
       # If we pull in the subsystem_rule() as well from this file, we get an error saying the scope
       # 'download-pex-bin' was not found when trying to fetch the appropriate scope.
       download_pex_bin.download_pex_bin,
-      *inject_init.rules(),
       *pex.rules(),
       *pex_from_target_closure.rules(),
       *prepare_chrooted_python_sources.rules(),

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

@@ -6,7 +6,6 @@
 from pants.backend.python.python_requirements import PythonRequirements
 from pants.backend.python.rules import (
   download_pex_bin,
-  inject_init,
   pex,
   pex_from_target_closure,
   prepare_chrooted_python_sources,
@@ -96,7 +95,6 @@ def register_goals():
 def rules():
   return (
     *download_pex_bin.rules(),
-    *inject_init.rules(),
     *prepare_chrooted_python_sources.rules(),
     *pex.rules(),
     *pex_from_target_closure.rules(),

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

@@ -41,6 +41,7 @@ async def create_pex_from_target_closure(request: CreatePexFromTargetClosure,
     python_setup=python_setup
   )
 
+  chrooted_sources: Optional[ChrootedPythonSources] = None
   if request.include_source_files:
     chrooted_sources = await Get[ChrootedPythonSources](HydratedTargets(all_targets))
 
@@ -54,7 +55,9 @@ async def create_pex_from_target_closure(request: CreatePexFromTargetClosure,
     requirements=requirements,
     interpreter_constraints=interpreter_constraints,
     entry_point=request.entry_point,
-    input_files_digest=chrooted_sources.digest if request.include_source_files else None,
+    input_files_digest=(
+      chrooted_sources.snapshot.directory_digest if chrooted_sources else None
+    ),
     additional_args=request.additional_args,
   )
 

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

@@ -3,47 +3,62 @@
 
 from dataclasses import dataclass
 
-from pants.backend.python.rules.inject_init import InjectedInitDigest
-from pants.engine.fs import Digest, DirectoriesToMerge
+from pants.engine.fs import (
+  EMPTY_DIRECTORY_DIGEST,
+  Digest,
+  DirectoriesToMerge,
+  FileContent,
+  InputFilesContent,
+  Snapshot,
+)
 from pants.engine.legacy.graph import HydratedTarget, HydratedTargets
 from pants.engine.rules import rule
 from pants.engine.selectors import Get, MultiGet
+from pants.python.pex_build_util import identify_missing_init_files
 from pants.rules.core.strip_source_roots import SourceRootStrippedSources
 
 
 @dataclass(frozen=True)
 class ChrootedPythonSources:
-  digest: Digest
+  snapshot: Snapshot
 
 
 @rule
-async def prepare_chrooted_python_sources(hydrated_targets: HydratedTargets) -> ChrootedPythonSources:
-  """Prepares Python sources by stripping the source root and injecting missing init.py files.
+async def prepare_chrooted_python_sources(
+  hydrated_targets: HydratedTargets,
+) -> ChrootedPythonSources:
+  """Prepares Python sources by stripping the source root and injecting missing __init__.py files.
 
   NB: This is useful for Pytest or ./pants run, but not every Python rule will need this.
   For example, autoformatters like Black do not need to understand relative imports or
   execute the code, so they can safely operate on the original source files without
   stripping source roots.
   """
-
   source_root_stripped_sources = await MultiGet(
     Get[SourceRootStrippedSources](HydratedTarget, hydrated_target)
     for hydrated_target in hydrated_targets
   )
+  sources_snapshot = await Get[Snapshot](
+    DirectoriesToMerge(
+      directories=tuple(
+        stripped_sources.snapshot.directory_digest
+        for stripped_sources in source_root_stripped_sources
+      )
+    )
+  )
 
-  sources_digest = await Get[Digest](DirectoriesToMerge(
-    directories=tuple(
-      stripped_sources.snapshot.directory_digest for stripped_sources in source_root_stripped_sources
+  missing_init_files = sorted(identify_missing_init_files(sources_snapshot.files))
+  inits_digest = EMPTY_DIRECTORY_DIGEST
+  if missing_init_files:
+    inits_digest = await Get[Digest](
+      InputFilesContent(FileContent(path=fp, content=b"") for fp in missing_init_files)
     )
-  ))
-  inits_digest = await Get[InjectedInitDigest](Digest, sources_digest)
-  sources_digest = await Get[Digest](DirectoriesToMerge(
-    directories=(sources_digest, inits_digest.directory_digest)
-  ))
-  return ChrootedPythonSources(digest=sources_digest)
+
+  result = await Get[Snapshot](
+    DirectoriesToMerge(directories=(sources_snapshot.directory_digest, inits_digest))
+  )
+  return ChrootedPythonSources(result)
 
 
 def rules():
-  return [
-    prepare_chrooted_python_sources,
-  ]
+  return [prepare_chrooted_python_sources]

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

@@ -0,0 +1,67 @@
+# Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from pathlib import PurePath
+from typing import List, Optional
+from unittest.mock import Mock
+
+from pants.backend.python.rules.prepare_chrooted_python_sources import (
+  ChrootedPythonSources,
+  prepare_chrooted_python_sources,
+)
+from pants.build_graph.address import Address
+from pants.build_graph.files import Files
+from pants.engine.legacy.graph import HydratedTarget, HydratedTargets
+from pants.engine.rules import RootRule
+from pants.engine.selectors import Params
+from pants.rules.core import strip_source_roots
+from pants.testutil.option.util import create_options_bootstrapper
+from pants.testutil.test_base import TestBase
+
+
+class PrepareChrootedPythonSourcesTest(TestBase):
+
+  @classmethod
+  def rules(cls):
+    return (
+      *super().rules(),
+      *strip_source_roots.rules(),
+      prepare_chrooted_python_sources,
+      RootRule(HydratedTargets),
+    )
+
+  def make_hydrated_target(
+    self, *, source_paths: List[str], type_alias: Optional[str] = None,
+  ) -> HydratedTarget:
+    adaptor = Mock()
+    adaptor.type_alias = type_alias
+    adaptor.sources = Mock()
+    adaptor.sources.snapshot = self.make_snapshot({fp: "" for fp in source_paths})
+    address = Address(spec_path=PurePath(source_paths[0]).parent.as_posix(), target_name="target")
+    return HydratedTarget(address=address, adaptor=adaptor, dependencies=())
+
+  def test_adds_missing_inits_and_strips_source_roots(self) -> None:
+    target_with_init = self.make_hydrated_target(
+      source_paths=["src/python/project/lib.py", "src/python/project/__init__.py"],
+    )
+    target_without_init = self.make_hydrated_target(
+      source_paths=["tests/python/test_project/f1.py", "tests/python/test_project/f2.py"],
+    )
+    files_target = self.make_hydrated_target(
+      source_paths=["src/python/project/resources/loose_file.txt"], type_alias=Files.alias(),
+    )
+    result = self.request_single_product(
+      ChrootedPythonSources,
+      Params(
+        HydratedTargets([target_with_init, target_without_init, files_target]),
+        create_options_bootstrapper(),
+      ),
+    )
+    assert sorted(result.snapshot.files) == sorted([
+      "project/lib.py",
+      "project/__init__.py",
+      "test_project/f1.py",
+      "test_project/f2.py",
+      "test_project/__init__.py",
+      "src/python/project/resources/loose_file.txt"
+    ])

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

@@ -122,7 +122,7 @@ async def setup_pytest_for_target(
 
   chrooted_sources = await Get[ChrootedPythonSources](HydratedTargets(all_targets))
   directories_to_merge = [
-    chrooted_sources.digest,
+    chrooted_sources.snapshot.directory_digest,
     resolved_requirements_pex.directory_digest,
   ]
 

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

@@ -8,7 +8,6 @@
 
 from pants.backend.python.rules import (
   download_pex_bin,
-  inject_init,
   pex,
   pex_from_target_closure,
   prepare_chrooted_python_sources,
@@ -109,7 +108,6 @@ def rules(cls):
       *super().rules(),
       *python_test_runner.rules(),
       *download_pex_bin.rules(),
-      *inject_init.rules(),
       *pex.rules(),
       *pex_from_target_closure.rules(),
       *prepare_chrooted_python_sources.rules(),