[internal] Refactor setup of GOROOT and import_analysis.py (#13000)

Prework for #12772.

src/python/pants/backend/go/lint/gofmt/rules.py

@@ -1,21 +1,23 @@
 # Copyright 2021 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
+from __future__ import annotations
+
 import dataclasses
+import os.path
 from dataclasses import dataclass
 
 from pants.backend.go.lint.fmt import GoLangFmtRequest
 from pants.backend.go.lint.gofmt.skip_field import SkipGofmtField
 from pants.backend.go.lint.gofmt.subsystem import GofmtSubsystem
-from pants.backend.go.subsystems.golang import GoLangDistribution
+from pants.backend.go.subsystems import golang
+from pants.backend.go.subsystems.golang import GoRoot
 from pants.backend.go.target_types import GoSources
 from pants.core.goals.fmt import FmtResult
 from pants.core.goals.lint import LintRequest, LintResult, LintResults
-from pants.core.util_rules.external_tool import DownloadedExternalTool, ExternalToolRequest
 from pants.core.util_rules.source_files import SourceFiles, SourceFilesRequest
 from pants.engine.fs import Digest, MergeDigests
-from pants.engine.internals.selectors import Get, MultiGet
-from pants.engine.platform import Platform
+from pants.engine.internals.selectors import Get
 from pants.engine.process import FallibleProcessResult, Process, ProcessResult
 from pants.engine.rules import collect_rules, rule
 from pants.engine.target import FieldSet, Target
@@ -52,20 +54,11 @@ class Setup:
 
 
 @rule(level=LogLevel.DEBUG)
-async def setup_gofmt(setup_request: SetupRequest, goroot: GoLangDistribution) -> Setup:
-    download_goroot_request = Get(
-        DownloadedExternalTool,
-        ExternalToolRequest,
-        goroot.get_request(Platform.current),
-    )
-
-    source_files_request = Get(
+async def setup_gofmt(setup_request: SetupRequest, goroot: GoRoot) -> Setup:
+    source_files = await Get(
         SourceFiles,
         SourceFilesRequest(field_set.sources for field_set in setup_request.request.field_sets),
     )
-
-    downloaded_goroot, source_files = await MultiGet(download_goroot_request, source_files_request)
-
     source_files_snapshot = (
         source_files.snapshot
         if setup_request.request.prior_formatter_result is None
@@ -74,14 +67,13 @@ async def setup_gofmt(setup_request: SetupRequest, goroot: GoLangDistribution) -
 
     input_digest = await Get(
         Digest,
-        MergeDigests((source_files_snapshot.digest, downloaded_goroot.digest)),
+        MergeDigests((source_files_snapshot.digest, goroot.digest)),
     )
-
-    argv = [
-        "./go/bin/gofmt",
+    argv = (
+        os.path.join(goroot.path, "bin/gofmt"),
         "-l" if setup_request.check_only else "-w",
         *source_files_snapshot.files,
-    ]
+    )
 
     process = Process(
         argv=argv,
@@ -90,7 +82,6 @@ async def setup_gofmt(setup_request: SetupRequest, goroot: GoLangDistribution) -
         description=f"Run gofmt on {pluralize(len(source_files_snapshot.files), 'file')}.",
         level=LogLevel.DEBUG,
     )
-
     return Setup(process=process, original_digest=source_files_snapshot.digest)
 
 
@@ -126,6 +117,7 @@ async def gofmt_lint(request: GofmtRequest, gofmt: GofmtSubsystem) -> LintResult
 def rules():
     return [
         *collect_rules(),
+        *golang.rules(),
         UnionRule(GoLangFmtRequest, GofmtRequest),
         UnionRule(LintRequest, GofmtRequest),
     ]

src/python/pants/backend/go/subsystems/golang.py

@@ -1,12 +1,21 @@
 # Copyright 2021 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
-from pants.core.util_rules.external_tool import TemplatedExternalTool
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from pants.core.util_rules.external_tool import (
+    DownloadedExternalTool,
+    ExternalToolRequest,
+    TemplatedExternalTool,
+)
+from pants.engine.fs import Digest
 from pants.engine.platform import Platform
-from pants.engine.rules import collect_rules
+from pants.engine.rules import Get, collect_rules, rule
 
 
-class GoLangDistribution(TemplatedExternalTool):
+class GolangSubsystem(TemplatedExternalTool):
     options_scope = "golang"
     name = "golang"
     help = "Official golang distribution."
@@ -28,5 +37,23 @@ def generate_exe(self, plat: Platform) -> str:
         return "./bin"
 
 
+@dataclass(frozen=True)
+class GoRoot:
+    """Path to the Go installation (the `GOROOT`)."""
+
+    path: str
+    digest: Digest
+
+
+@rule
+async def setup_goroot(golang_subsystem: GolangSubsystem) -> GoRoot:
+    downloaded_go_dist = await Get(
+        DownloadedExternalTool,
+        ExternalToolRequest,
+        golang_subsystem.get_request(Platform.current),
+    )
+    return GoRoot("./go", downloaded_go_dist.digest)
+
+
 def rules():
     return collect_rules()

src/python/pants/backend/go/target_type_rules.py

@@ -31,7 +31,7 @@
     ResolveGoModuleRequest,
 )
 from pants.backend.go.util_rules.go_pkg import ResolvedGoPackage, ResolveGoPackageRequest
-from pants.backend.go.util_rules.import_analysis import ResolvedImportPathsForGoLangDistribution
+from pants.backend.go.util_rules.import_analysis import GoStdLibImports
 from pants.base.specs import (
     AddressSpecs,
     DescendantAddresses,
@@ -113,7 +113,7 @@ class InferGoPackageDependenciesRequest(InferDependenciesRequest):
 @rule
 async def infer_go_dependencies(
     request: InferGoPackageDependenciesRequest,
-    goroot_imports: ResolvedImportPathsForGoLangDistribution,
+    std_lib_imports: GoStdLibImports,
     package_mapping: GoImportPathToPackageMapping,
 ) -> InferredDependencies:
     this_go_package = await Get(
@@ -156,8 +156,7 @@ async def infer_go_dependencies(
     # external modules.
     inferred_dependencies = []
     for import_path in this_go_package.imports + this_go_package.test_imports:
-        # Check whether the import path comes from the standard library.
-        if import_path in goroot_imports.import_path_mapping:
+        if import_path in std_lib_imports:
             continue
 
         # Infer first-party dependencies to other packages in same go_module.
@@ -194,7 +193,7 @@ class InjectGoExternalPackageDependenciesRequest(InjectDependenciesRequest):
 @rule
 async def inject_go_external_package_dependencies(
     request: InjectGoExternalPackageDependenciesRequest,
-    goroot_imports: ResolvedImportPathsForGoLangDistribution,
+    std_lib_imports: GoStdLibImports,
     package_mapping: GoImportPathToPackageMapping,
 ) -> InjectedDependencies:
     this_go_package = await Get(
@@ -205,8 +204,7 @@ async def inject_go_external_package_dependencies(
     # external modules.
     inferred_dependencies = []
     for import_path in this_go_package.imports + this_go_package.test_imports:
-        # Check whether the import path comes from the standard library.
-        if import_path in goroot_imports.import_path_mapping:
+        if import_path in std_lib_imports:
             continue
 
         # Infer third-party dependencies on _go_external_package targets.

src/python/pants/backend/go/util_rules/import_analysis.py

@@ -3,21 +3,18 @@
 
 from __future__ import annotations
 
-import json
 import logging
 import os
-import textwrap
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Dict, List
+from typing import TYPE_CHECKING
 
 import ijson
 
-from pants.backend.go.subsystems.golang import GoLangDistribution
-from pants.core.util_rules.external_tool import DownloadedExternalTool, ExternalToolRequest
+from pants.backend.go.subsystems.golang import GoRoot
+from pants.backend.go.util_rules.sdk import GoSdkProcess
 from pants.engine.fs import AddPrefix, CreateDigest, Digest, FileContent, MergeDigests
 from pants.engine.internals.selectors import Get
-from pants.engine.platform import Platform
-from pants.engine.process import BashBinary, Process, ProcessResult
+from pants.engine.process import ProcessResult
 from pants.engine.rules import collect_rules, rule
 from pants.util.frozendict import FrozenDict
 from pants.util.logging import LogLevel
@@ -30,94 +27,32 @@
 logger = logging.getLogger(__name__)
 
 
-@dataclass(frozen=True)
-class ImportDescriptor:
-    digest: Digest
-    path: str
-
-
-# Note: There is only one subclass of this class currently. There will be additional subclasses once module
-# support is added to the plugin.
-@dataclass(frozen=True)
-class ResolvedImportPaths:
-    """Base class for types which map import paths provided by a source to how to access built
-    package files for those import paths."""
-
-    import_path_mapping: FrozenDict[str, ImportDescriptor]
-
-
-@dataclass(frozen=True)
-class ResolvedImportPathsForGoLangDistribution(ResolvedImportPaths):
-    pass
-
+class GoStdLibImports(FrozenDict[str, str]):
+    """A mapping of standard library import paths to the `.a` static file paths for that import
+    path.
 
-def parse_imports_for_golang_distribution(raw_json: bytes) -> Dict[str, str]:
-    import_paths: Dict[str, str] = {}
-    package_descriptors = ijson.items(raw_json, "", multiple_values=True)
-    for package_descriptor in package_descriptors:
-        try:
-            if "Target" in package_descriptor and "ImportPath" in package_descriptor:
-                import_paths[package_descriptor["ImportPath"]] = package_descriptor["Target"]
-        except Exception as ex:
-            logger.error(
-                f"error while parsing package descriptor: {ex}; package_descriptor: {json.dumps(package_descriptor)}"
-            )
-            raise
-    return import_paths
+    For example, "net/smtp": "/absolute_path_to_goroot/pkg/darwin_arm64/net/smtp.a".
+    """
 
 
-@rule
-async def analyze_imports_for_golang_distribution(
-    goroot: GoLangDistribution,
-    platform: Platform,
-    bash: BashBinary,
-) -> ResolvedImportPathsForGoLangDistribution:
-    downloaded_goroot = await Get(
-        DownloadedExternalTool,
-        ExternalToolRequest,
-        goroot.get_request(platform),
-    )
-
-    # Note: The `go` tool requires GOPATH to be an absolute path which can only be resolved from within the
-    # execution sandbox. Thus, this code uses a bash script to be able to resolve that path.
-    analyze_script_digest = await Get(
-        Digest,
-        CreateDigest(
-            [
-                FileContent(
-                    "analyze.sh",
-                    textwrap.dedent(
-                        """\
-                export GOROOT="./go"
-                export GOPATH="$(/bin/pwd)/gopath"
-                export GOCACHE="$(/bin/pwd)/cache"
-                mkdir -p "$GOPATH" "$GOCACHE"
-                exec ./go/bin/go list -json std
-                """
-                    ).encode("utf-8"),
-                )
-            ]
+@rule(desc="Determine Go std lib's imports", level=LogLevel.DEBUG)
+async def determine_go_std_lib_imports() -> GoStdLibImports:
+    list_result = await Get(
+        ProcessResult,
+        GoSdkProcess(
+            command=("list", "-json", "std"),
+            description="Ask Go for its available import paths",
+            absolutify_goroot=False,
         ),
     )
-
-    input_root = await Get(Digest, MergeDigests([downloaded_goroot.digest, analyze_script_digest]))
-
-    process = Process(
-        argv=[bash.path, "./analyze.sh"],
-        input_digest=input_root,
-        description="Analyze import paths available in Go distribution.",
-        level=LogLevel.DEBUG,
-    )
-
-    result = await Get(ProcessResult, Process, process)
-    import_paths = parse_imports_for_golang_distribution(result.stdout)
-    import_descriptors: Dict[str, ImportDescriptor] = {
-        import_path: ImportDescriptor(digest=downloaded_goroot.digest, path=path)
-        for import_path, path in import_paths.items()
-    }
-    return ResolvedImportPathsForGoLangDistribution(
-        import_path_mapping=FrozenDict(import_descriptors)
-    )
+    result = {}
+    for package_descriptor in ijson.items(list_result.stdout, "", multiple_values=True):
+        import_path = package_descriptor.get("ImportPath")
+        target = package_descriptor.get("Target")
+        if not import_path or not target:
+            continue
+        result[import_path] = target
+    return GoStdLibImports(result)
 
 
 @dataclass(frozen=True)
@@ -133,7 +68,7 @@ class GatheredImports:
 
 @rule
 async def generate_import_config(
-    request: GatherImportsRequest, goroot_import_mappings: ResolvedImportPathsForGoLangDistribution
+    request: GatherImportsRequest, stdlib_imports: GoStdLibImports, goroot: GoRoot
 ) -> GatheredImports:
     import_config_digests: dict[str, tuple[str, Digest]] = {}
     for pkg in request.packages:
@@ -143,27 +78,25 @@ async def generate_import_config(
 
     pkg_digests: OrderedSet[Digest] = OrderedSet()
 
-    import_config: List[str] = ["# import config"]
+    import_config = ["# import config"]
     for import_path, (fp, digest) in import_config_digests.items():
         pkg_digests.add(digest)
         import_config.append(f"packagefile {import_path}=__pkgs__/{fp}/__pkg__.a")
 
     if request.include_stdlib:
-        for stdlib_pkg_importpath, stdlib_pkg in goroot_import_mappings.import_path_mapping.items():
-            pkg_digests.add(stdlib_pkg.digest)
-            import_config.append(
-                f"packagefile {stdlib_pkg_importpath}={os.path.normpath(stdlib_pkg.path)}"
-            )
+        pkg_digests.add(goroot.digest)
+        import_config.extend(
+            f"packagefile {import_path}={os.path.normpath(static_file_path)}"
+            for import_path, static_file_path in stdlib_imports.items()
+        )
 
     import_config_content = "\n".join(import_config).encode("utf-8")
-
     import_config_digest = await Get(
-        Digest, CreateDigest([FileContent(path="./importcfg", content=import_config_content)])
+        Digest, CreateDigest([FileContent("./importcfg", import_config_content)])
     )
     pkg_digests.add(import_config_digest)
 
     digest = await Get(Digest, MergeDigests(pkg_digests))
-
     return GatheredImports(digest=digest)