Properly handle PEP 420 namespace packages

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

@@ -1,21 +1,16 @@
 # Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
-import os
+import itertools
 from dataclasses import dataclass
+from pathlib import PurePath
 
-from pants.engine.fs import (
-    Digest,
-    FileContent,
-    InputFilesContent,
-    MergeDigests,
-    PathGlobs,
-    Snapshot,
-)
-from pants.engine.rules import RootRule, rule
-from pants.engine.selectors import Get, MultiGet
+from pants.core.util_rules.strip_source_roots import SourceRootStrippedSources, StripSnapshotRequest
+from pants.engine.fs import MergeDigests, PathGlobs, Snapshot
+from pants.engine.rules import rule
+from pants.engine.selectors import Get
 from pants.python.pex_build_util import identify_missing_init_files
-from pants.source.source_root import OptionalSourceRoot, SourceRootRequest
+from pants.source.source_root import AllSourceRoots
 from pants.util.ordered_set import FrozenOrderedSet
 
 
@@ -31,52 +26,36 @@ class InitInjectedSnapshot:
 
 
 @rule
-async def inject_missing_init_files(request: InjectInitRequest) -> InitInjectedSnapshot:
-    """Ensure that every package has an `__init__.py` file in it.
-
-    This will first use any `__init__.py` files in the input snapshot, then read from the filesystem
-    to see if any exist but are not in the snapshot, and finally will create empty files.
-    """
-    original_missing_init_files = identify_missing_init_files(request.snapshot.files)
-    if not original_missing_init_files:
+async def inject_missing_init_files(
+    request: InjectInitRequest, all_source_roots: AllSourceRoots
+) -> InitInjectedSnapshot:
+    """Add any `__init__.py` files that exist on the filesystem but are not yet in the snapshot."""
+    missing_init_files = identify_missing_init_files(request.snapshot.files)
+    if not missing_init_files:
         return InitInjectedSnapshot(request.snapshot)
 
-    missing_init_files = original_missing_init_files
-    if not request.sources_stripped:
-        # Get rid of any identified-as-missing __init__.py files that are not under a source root.
-        optional_src_roots = await MultiGet(
-            Get(OptionalSourceRoot, SourceRootRequest, SourceRootRequest.for_file(init_file))
-            for init_file in original_missing_init_files
-        )
-
-        def is_under_source_root(init_file: str, optional_src_root: OptionalSourceRoot) -> bool:
-            return (
-                optional_src_root.source_root is not None
-                and optional_src_root.source_root.path != os.path.dirname(init_file)
-            )
-
+    if request.sources_stripped:
+        # If files are stripped, we don't know what source root they might live in, so we look
+        # up every source root.
+        roots = tuple(root.path for root in all_source_roots)
         missing_init_files = FrozenOrderedSet(
-            init_file
-            for init_file, optional_src_root in zip(original_missing_init_files, optional_src_roots)
-            if is_under_source_root(init_file, optional_src_root)
+            PurePath(root, f).as_posix() for root, f in itertools.product(roots, missing_init_files)
         )
 
+    # NB: This will intentionally _not_ error on any unmatched globs.
     discovered_inits_snapshot = await Get(Snapshot, PathGlobs(missing_init_files))
-    generated_inits_digest = await Get(
-        Digest,
-        InputFilesContent(
-            FileContent(fp, b"# Generated by Pants to ensure that this package is importable.")
-            for fp in missing_init_files.difference(discovered_inits_snapshot.files)
-        ),
-    )
+    if request.sources_stripped:
+        # We must now strip all discovered paths.
+        stripped_snapshot = await Get(
+            SourceRootStrippedSources, StripSnapshotRequest(discovered_inits_snapshot)
+        )
+        discovered_inits_snapshot = stripped_snapshot.snapshot
+
     result = await Get(
-        Snapshot,
-        MergeDigests(
-            (request.snapshot.digest, discovered_inits_snapshot.digest, generated_inits_digest)
-        ),
+        Snapshot, MergeDigests((request.snapshot.digest, discovered_inits_snapshot.digest)),
     )
     return InitInjectedSnapshot(result)
 
 
 def rules():
-    return [inject_missing_init_files, RootRule(InjectInitRequest)]
+    return [inject_missing_init_files]

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

@@ -1,14 +1,15 @@
 # Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
-from typing import List, Optional
+from typing import List
 
 from pants.backend.python.rules.inject_init import (
     InitInjectedSnapshot,
     InjectInitRequest,
     inject_missing_init_files,
 )
-from pants.engine.fs import Digest, FilesContent
+from pants.core.util_rules import strip_source_roots
+from pants.engine.fs import FilesContent
 from pants.engine.rules import RootRule
 from pants.engine.selectors import Params
 from pants.testutil.option.util import create_options_bootstrapper
@@ -21,90 +22,68 @@ def rules(cls):
         return (
             *super().rules(),
             inject_missing_init_files,
+            *strip_source_roots.rules(),
             RootRule(InjectInitRequest),
-            RootRule(Digest),
         )
 
     def assert_injected(
         self,
         *,
-        original_files: List[str],
-        expected_added: List[str],
-        expected_discovered: Optional[List[str]] = None,
-        sources_stripped=True,
+        declared_files_stripped: bool,
+        original_declared_files: List[str],
+        original_undeclared_files: List[str],
+        expected_discovered: List[str],
     ) -> None:
-        expected_discovered = expected_discovered or []
+        for f in original_undeclared_files:
+            self.create_file(f, "# undeclared")
         request = InjectInitRequest(
-            self.make_snapshot({fp: "# python code" for fp in original_files}),
-            sources_stripped=sources_stripped,
+            self.make_snapshot({fp: "# declared" for fp in original_declared_files}),
+            sources_stripped=declared_files_stripped,
         )
         result = self.request_single_product(
-            InitInjectedSnapshot,
-            Params(
-                request, create_options_bootstrapper(args=["--source-root-patterns=['src/python']"])
-            ),
+            InitInjectedSnapshot, Params(request, create_options_bootstrapper())
         ).snapshot
-        assert list(result.files) == sorted(
-            [*original_files, *expected_added, *expected_discovered]
-        )
-        # Ensure all original `__init__.py` are preserved with their original content.
-        materialized_original_inits = [
-            fc
-            for fc in self.request_single_product(FilesContent, result.digest)
-            if fc.path.endswith("__init__.py")
-            and (fc.path in original_files or fc.path in expected_discovered)
-        ]
-        for original_init in materialized_original_inits:
-            assert (
-                original_init.content == b"# python code"
-            ), f"{original_init} does not have its original content preserved."
+        assert list(result.files) == sorted([*original_declared_files, *expected_discovered])
 
-    def test_no_inits_present(self) -> None:
-        self.assert_injected(
-            original_files=["lib.py", "subdir/lib.py"], expected_added=["subdir/__init__.py"],
-        )
-        self.assert_injected(
-            original_files=["a/b/lib.py", "a/b/subdir/lib.py"],
-            expected_added=["a/__init__.py", "a/b/__init__.py", "a/b/subdir/__init__.py",],
-        )
+        materialized_result = self.request_single_product(FilesContent, result.digest)
+        for file_content in materialized_result:
+            path = file_content.path
+            if not path.endswith("__init__.py"):
+                continue
+            assert path in original_declared_files or path in expected_discovered
+            expected = b"# declared" if path in original_declared_files else b"# undeclared"
+            assert file_content.content == expected
 
-    def test_preserves_original_inits(self) -> None:
-        self.assert_injected(
-            original_files=["lib.py", "__init__.py", "subdir/lib.py"],
-            expected_added=["subdir/__init__.py"],
-        )
+    def test_unstripped(self) -> None:
         self.assert_injected(
-            original_files=[
-                "a/b/lib.py",
-                "a/b/__init__.py",
-                "a/b/subdir/lib.py",
-                "a/b/subdir/__init__.py",
+            declared_files_stripped=False,
+            original_declared_files=[
+                "src/python/project/lib.py",
+                "src/python/project/subdir/__init__.py",
+                "src/python/project/subdir/lib.py",
+                "src/python/no_init/lib.py",
             ],
-            expected_added=["a/__init__.py"],
-        )
-        # No missing `__init__.py` files
-        self.assert_injected(
-            original_files=["lib.py", "__init__.py", "subdir/lib.py", "subdir/__init__.py"],
-            expected_added=[],
-        )
-
-    def test_finds_undeclared_original_inits(self) -> None:
-        self.create_file("a/__init__.py", "# python code")
-        self.create_file("a/b/__init__.py", "# python code")
-        self.assert_injected(
-            original_files=["a/b/subdir/lib.py"],
-            expected_added=["a/b/subdir/__init__.py"],
-            expected_discovered=["a/__init__.py", "a/b/__init__.py"],
+            original_undeclared_files=[
+                "src/python/project/__init__.py",
+                "tests/python/project/__init__.py",
+            ],
+            expected_discovered=["src/python/project/__init__.py"],
         )
 
-    def test_source_roots_unstripped(self) -> None:
+    def test_stripped(self) -> None:
         self.assert_injected(
-            original_files=[
-                "src/python/lib.py",
-                "src/python/subdir/lib.py",
-                "src/python/subdir/__init__.py",
-                "src/python/another_subdir/lib.py",
+            declared_files_stripped=True,
+            original_declared_files=[
+                "project/lib.py",
+                "project/subdir/lib.py",
+                "project/subdir/__init__.py",
+                "project/no_init/lib.py",
+            ],
+            # NB: These will strip down to end up being the same file. If they had different
+            # contents, Pants would error when trying to merge them.
+            original_undeclared_files=[
+                "src/python/project/__init__.py",
+                "tests/python/project/__init__.py",
             ],
-            expected_added=["src/python/another_subdir/__init__.py"],
-            sources_stripped=False,
+            expected_discovered=["project/__init__.py"],
         )

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

@@ -123,6 +123,9 @@ async def prepare_unstripped_python_sources(
             strip_source_roots=False,
         ),
     )
+    init_injected = await Get(
+        InitInjectedSnapshot, InjectInitRequest(sources.snapshot, sources_stripped=False)
+    )
 
     source_root_objs = await MultiGet(
         Get(
@@ -134,10 +137,6 @@ async def prepare_unstripped_python_sources(
         if tgt.has_field(PythonSources) or tgt.has_field(ResourcesSources)
     )
     source_root_paths = {source_root_obj.path for source_root_obj in source_root_objs}
-
-    init_injected = await Get(
-        InitInjectedSnapshot, InjectInitRequest(sources.snapshot, sources_stripped=False)
-    )
     return UnstrippedPythonSources(init_injected.snapshot, tuple(sorted(source_root_paths)))