javascript: Support Node.js subpath imports (pantsbuild#18934)

Adds support for the inference impl to inspect `package.json#imports`,
and backwards map this object to find dependencies known by pants.

Relevant part of Node.js docs:
https://nodejs.org/api/packages.html#subpath-imports.

src/python/pants/backend/javascript/dependency_inference/rules.py

@@ -22,6 +22,7 @@
     OwningNodePackage,
     OwningNodePackageRequest,
     PackageJsonEntryPoints,
+    PackageJsonImports,
     PackageJsonSourceField,
 )
 from pants.backend.javascript.subsystems.nodejs_infer import NodeJSInfer
@@ -98,6 +99,22 @@ async def map_candidate_node_packages(
     )
 
 
+async def _replace_subpath_imports(
+    req: InferJSDependenciesRequest, import_strings: JSImportStrings
+) -> JSImportStrings:
+    owning_pkg = await Get(OwningNodePackage, OwningNodePackageRequest(req.field_set.address))
+    if owning_pkg.target:
+        subpath_imports = await Get(
+            PackageJsonImports, PackageJsonSourceField, owning_pkg.target[PackageJsonSourceField]
+        )
+        return JSImportStrings(
+            replace_string
+            for string in import_strings
+            for replace_string in subpath_imports.replacements(string) or (string,)
+        )
+    return import_strings
+
+
 @rule
 async def infer_js_source_dependencies(
     request: InferJSDependenciesRequest,
@@ -108,6 +125,8 @@ async def infer_js_source_dependencies(
         return InferredDependencies(())
 
     import_strings = await Get(JSImportStrings, ParseJsImportStrings(source))
+    import_strings = await _replace_subpath_imports(request, import_strings)
+
     path_strings = FrozenOrderedSet(
         os.path.normpath(os.path.join(os.path.dirname(source.file_path), import_string))
         for import_string in import_strings
@@ -118,8 +137,10 @@ async def infer_js_source_dependencies(
     owning_targets = await Get(Targets, Addresses(owners))
 
     non_path_string_bases = FrozenOrderedSet(
-        os.path.basename(non_path_string) for non_path_string in import_strings - path_strings
+        non_path_string.partition(os.path.sep)[0]
+        for non_path_string in import_strings - path_strings
     )
+
     candidate_pkgs = await Get(
         NodePackageCandidateMap, RequestNodePackagesCandidateMap(request.field_set.address)
     )

src/python/pants/backend/javascript/dependency_inference/rules_test.py

@@ -287,6 +287,68 @@ def test_infers_third_party_package_json_field_js_source_dependency(
     assert set(addresses) == {Address("src/js", generated_name="chalk")}
 
 
+def test_infers_third_party_package_json_field_js_source_dependency_with_import_subpaths(
+    rule_runner: RuleRunner,
+) -> None:
+    rule_runner.write_files(
+        {
+            "src/js/BUILD": "package_json()",
+            "src/js/package.json": given_package(
+                "ham",
+                "0.0.1",
+                main="lib/index.js",
+                dependencies={"chalk": "5.0.2"},
+                imports={"#myChalk": "chalk"},
+            ),
+            "src/js/lib/BUILD": "javascript_sources()",
+            "src/js/lib/index.js": dedent(
+                """\
+                import chalk from "#myChalk";
+                """
+            ),
+        }
+    )
+
+    pkg_tgt = rule_runner.get_target(Address("src/js/lib", relative_file_path="index.js"))
+    addresses = rule_runner.request(
+        InferredDependencies,
+        [InferJSDependenciesRequest(JSSourceInferenceFieldSet.create(pkg_tgt))],
+    ).include
+
+    assert set(addresses) == {Address("src/js", generated_name="chalk")}
+
+
+def test_infers_third_party_package_json_field_js_source_dependency_with_import_subpaths_with_star_replacements(
+    rule_runner: RuleRunner,
+) -> None:
+    rule_runner.write_files(
+        {
+            "src/js/BUILD": "package_json()",
+            "src/js/package.json": given_package(
+                "ham",
+                "0.0.1",
+                main="lib/index.js",
+                dependencies={"chalk": "5.0.2"},
+                imports={"#myChalk/*.js": "chalk/stuff/*.js"},
+            ),
+            "src/js/lib/BUILD": "javascript_sources()",
+            "src/js/lib/index.js": dedent(
+                """\
+                import chalk from "#myChalk/index.js";
+                """
+            ),
+        }
+    )
+
+    pkg_tgt = rule_runner.get_target(Address("src/js/lib", relative_file_path="index.js"))
+    addresses = rule_runner.request(
+        InferredDependencies,
+        [InferJSDependenciesRequest(JSSourceInferenceFieldSet.create(pkg_tgt))],
+    ).include
+
+    assert set(addresses) == {Address("src/js", generated_name="chalk")}
+
+
 def test_infers_first_party_package_json_field_js_source_dependency(
     rule_runner: RuleRunner,
 ) -> None:
@@ -314,3 +376,63 @@ def test_infers_first_party_package_json_field_js_source_dependency(
     ).include
 
     assert set(addresses) == {Address("src/js/b", generated_name="spam")}
+
+
+def test_infers_first_party_package_json_field_js_source_dependency_with_import_subpaths(
+    rule_runner: RuleRunner,
+) -> None:
+    rule_runner.write_files(
+        {
+            "src/js/a/BUILD": "package_json()",
+            "src/js/a/package.json": given_package("ham", "0.0.1", imports={"#spam": "spam"}),
+            "src/js/a/lib/BUILD": "javascript_sources()",
+            "src/js/a/lib/index.js": dedent(
+                """\
+                import { x } from "#spam";
+                """
+            ),
+            "src/js/b/BUILD": "package_json()",
+            "src/js/b/package.json": given_package("spam", "0.0.1"),
+            "src/js/b/lib/BUILD": "javascript_sources()",
+            "src/js/b/lib/index.js": "const x = 2;",
+        }
+    )
+
+    pkg_tgt = rule_runner.get_target(Address("src/js/a/lib", relative_file_path="index.js"))
+    addresses = rule_runner.request(
+        InferredDependencies,
+        [InferJSDependenciesRequest(JSSourceInferenceFieldSet.create(pkg_tgt))],
+    ).include
+
+    assert set(addresses) == {Address("src/js/b", generated_name="spam")}
+
+
+def test_infers_first_party_package_json_field_js_source_dependency_with_starred_import_subpaths(
+    rule_runner: RuleRunner,
+) -> None:
+    rule_runner.write_files(
+        {
+            "src/js/a/BUILD": "package_json()",
+            "src/js/a/package.json": given_package(
+                "ham", "0.0.1", imports={"#spam/*.js": "spam/lib/*.js"}
+            ),
+            "src/js/a/lib/BUILD": "javascript_sources()",
+            "src/js/a/lib/index.js": dedent(
+                """\
+                import { x } from "#spam/index.js";
+                """
+            ),
+            "src/js/b/BUILD": "package_json()",
+            "src/js/b/package.json": given_package("spam", "0.0.1"),
+            "src/js/b/lib/BUILD": "javascript_sources()",
+            "src/js/b/lib/index.js": "const x = 2;",
+        }
+    )
+
+    pkg_tgt = rule_runner.get_target(Address("src/js/a/lib", relative_file_path="index.js"))
+    addresses = rule_runner.request(
+        InferredDependencies,
+        [InferJSDependenciesRequest(JSSourceInferenceFieldSet.create(pkg_tgt))],
+    ).include
+
+    assert set(addresses) == {Address("src/js/b", generated_name="spam")}

src/python/pants/backend/javascript/package_json.py

@@ -4,7 +4,9 @@
 
 import itertools
 import json
+import logging
 import os.path
+import re
 from abc import ABC
 from dataclasses import dataclass, field
 from typing import Any, ClassVar, Iterable, Mapping, Optional, Tuple
@@ -63,6 +65,8 @@
 from pants.util.frozendict import FrozenDict
 from pants.util.strutil import help_text, softwrap
 
+_logger = logging.getLogger(__name__)
+
 
 class NodePackageDependenciesField(Dependencies):
     pass
@@ -443,6 +447,83 @@ class NodeBuildScriptTarget(Target):
     )
 
 
+@dataclass(frozen=True)
+class PackageJsonImports:
+    """https://nodejs.org/api/packages.html#subpath-imports."""
+
+    imports: FrozenDict[re.Pattern[str], tuple[str, ...]]
+    root_dir: str
+
+    def replacements(self, import_string: str) -> tuple[str, ...]:
+        def replace_matching_pattern(
+            pattern: re.Pattern[str], subpath: str, string: str
+        ) -> str | None:
+            match = pattern.match(string)
+            if match:
+                replacement = subpath
+                for group in match.groups():
+                    replacement = replacement.replace("*", group, 1)
+                if "*" in replacement:
+                    _logger.warning(
+                        softwrap(
+                            f"""
+                            package.json#imports pattern '{pattern.pattern}' matched '{string}',
+                            but the resulting subpath '{subpath}' string replacements '*'
+                            did not match.
+
+                            Inference will not behave correctly for import '{string}'.
+                            """
+                        )
+                    )
+                    return None
+                return "".join((replacement, string[match.endpos :]))
+            return None
+
+        return tuple(
+            filter(
+                None,
+                (
+                    replace_matching_pattern(pattern, subpath, import_string)
+                    for pattern, subpaths in self.imports.items()
+                    for subpath in subpaths
+                ),
+            )
+        )
+
+    @classmethod
+    def from_package_json(cls, pkg_json: PackageJson) -> PackageJsonImports:
+        return cls(
+            imports=cls._import_from_package_json(pkg_json),
+            root_dir=pkg_json.root_dir,
+        )
+
+    @staticmethod
+    def _to_import_pattern(string: str) -> re.Pattern[str]:
+        return re.compile(r"^" + re.escape(string).replace(r"\*", "(.*)"))
+
+    @staticmethod
+    def _import_from_package_json(
+        pkg_json: PackageJson,
+    ) -> FrozenDict[re.Pattern[str], tuple[str, ...]]:
+        imports: Mapping[str, Any] | None = pkg_json.content.get("imports")
+
+        def get_subpaths(value: str | Mapping[str, Any]) -> Iterable[str]:
+            if isinstance(value, str):
+                yield value
+            elif isinstance(value, Mapping):
+                for v in value.values():
+                    yield from get_subpaths(v)
+
+        if not imports:
+            return FrozenDict()
+        return FrozenDict(
+            {
+                PackageJsonImports._to_import_pattern(key): tuple(sorted(get_subpaths(subpath)))
+                for key, subpath in imports.items()
+            }
+        )
+
+
 @dataclass(frozen=True)
 class PackageJsonEntryPoints:
     """See https://nodejs.org/api/packages.html#package-entry-points and
@@ -729,6 +810,15 @@ async def script_entrypoints_for_source(
     )
 
 
+@rule
+async def subpath_imports_for_source(
+    source_field: PackageJsonSourceField,
+) -> PackageJsonImports:
+    return PackageJsonImports.from_package_json(
+        await Get(PackageJson, PackageJsonSourceField, source_field)
+    )
+
+
 class GenerateNodePackageTargets(GenerateTargetsRequest):
     generate_from = PackageJsonTarget
 

src/python/pants/backend/javascript/package_json_test.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import json
+import re
 from textwrap import dedent
 from typing import Iterable
 
@@ -15,6 +16,8 @@
     NodeTestScript,
     NodeThirdPartyPackageTarget,
     PackageJson,
+    PackageJsonImports,
+    PackageJsonSourceField,
     PackageJsonTarget,
 )
 from pants.build_graph.address import Address
@@ -36,6 +39,7 @@ def rule_runner() -> RuleRunner:
             *package_json.rules(),
             QueryRule(AllPackageJson, ()),
             QueryRule(Owners, (OwnersRequest,)),
+            QueryRule(PackageJsonImports, (PackageJsonSourceField,)),
         ],
         target_types=[
             PackageJsonTarget,
@@ -269,3 +273,46 @@ def test_specifying_missing_custom_coverage_entry_point_script_is_an_error(
     )
     with pytest.raises(ExecutionError):
         rule_runner.get_target(Address("src/js", generated_name="ham"))
+
+
+def test_parses_subpath_imports(
+    rule_runner: RuleRunner,
+) -> None:
+    rule_runner.write_files(
+        {
+            "src/js/BUILD": dedent(
+                """\
+                package_json()
+                """
+            ),
+            "src/js/package.json": json.dumps(
+                {
+                    "name": "ham",
+                    "version": "0.0.1",
+                    "imports": {
+                        "#a": "./yep.js",
+                        "#b": "some-package",
+                        "#c": {"node": "polyfill", "default": "./polyfill.js"},
+                        "#d/module/js/*.js": "./module/*.js",
+                    },
+                }
+            ),
+        }
+    )
+
+    tgt = rule_runner.get_target(Address("src/js", generated_name="ham"))
+    imports = rule_runner.request(PackageJsonImports, (tgt[PackageJsonSourceField],))
+
+    assert imports.imports == FrozenDict(
+        {
+            re.compile(r"^\#a"): ("./yep.js",),  # noqa: W605 # Escape added by re.escape
+            re.compile(r"^\#b"): ("some-package",),  # noqa: W605 # Escape added by re.escape
+            re.compile("^\#c"): (  # noqa: W605 # Escape added by re.escape
+                "./polyfill.js",
+                "polyfill",
+            ),
+            re.compile(r"^\#d/module/js/(.*)\.js"): (  # noqa: W605 # Escape added by re.escape
+                "./module/*.js",
+            ),
+        }
+    )