[internal] refactor scrooge codegen to allow for multiple languages

pants.toml

@@ -27,7 +27,7 @@ backend_packages.add = [
   "pants.backend.experimental.scala",
   "pants.backend.experimental.scala.lint.scalafmt",
   "pants.backend.experimental.codegen.protobuf.scala",
-  "pants.backend.experimental.codegen.thrift.scrooge",
+  "pants.backend.experimental.codegen.thrift.scrooge.scala",
   "pants.backend.experimental.scala.debug_goals",
   "internal_plugins.releases",
 ]

src/python/pants/backend/codegen/thrift/scrooge/rules.py

@@ -1,74 +1,62 @@
 # Copyright 2021 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
+from dataclasses import dataclass
 
 from pants.backend.codegen.thrift.scrooge import additional_fields
 from pants.backend.codegen.thrift.scrooge.additional_fields import ScroogeFinagleBoolField
 from pants.backend.codegen.thrift.scrooge.subsystem import ScroogeSubsystem
-from pants.backend.codegen.thrift.target_types import ThriftDependenciesField, ThriftSourceField
-from pants.backend.java.target_types import JavaSourceField
+from pants.backend.codegen.thrift.target_types import ThriftSourceField
+from pants.build_graph.address import Address
 from pants.core.util_rules.source_files import SourceFiles, SourceFilesRequest
-from pants.engine.addresses import Addresses, UnparsedAddressInputs
-from pants.engine.fs import (
-    AddPrefix,
-    CreateDigest,
-    Digest,
-    Directory,
-    MergeDigests,
-    RemovePrefix,
-    Snapshot,
-)
+from pants.engine.fs import CreateDigest, Digest, Directory, MergeDigests, RemovePrefix, Snapshot
 from pants.engine.internals.selectors import Get, MultiGet
 from pants.engine.process import BashBinary, Process, ProcessResult
 from pants.engine.rules import collect_rules, rule
-from pants.engine.target import (
-    GeneratedSources,
-    GenerateSourcesRequest,
-    InjectDependenciesRequest,
-    InjectedDependencies,
-    TransitiveTargets,
-    TransitiveTargetsRequest,
-)
+from pants.engine.target import TransitiveTargets, TransitiveTargetsRequest, WrappedTarget
 from pants.engine.unions import UnionRule
 from pants.jvm.jdk_rules import JdkSetup
 from pants.jvm.resolve.coursier_fetch import MaterializedClasspath, MaterializedClasspathRequest
 from pants.jvm.resolve.jvm_tool import JvmToolLockfileRequest, JvmToolLockfileSentinel
-from pants.source.source_root import (
-    SourceRoot,
-    SourceRootRequest,
-    SourceRootsRequest,
-    SourceRootsResult,
-)
+from pants.source.source_root import SourceRootsRequest, SourceRootsResult
 from pants.util.logging import LogLevel
 
 
-class GenerateScalaFromThriftRequest(GenerateSourcesRequest):
-    input = ThriftSourceField
-    output = JavaSourceField
+@dataclass(frozen=True)
+class GenerateScroogeThriftSourcesRequest:
+    thrift_source_field: ThriftSourceField
+    lang_id: str
+    lang_name: str
+
+
+@dataclass(frozen=True)
+class GeneratedScroogeThriftSources:
+    snapshot: Snapshot
 
 
 class ScroogeToolLockfileSentinel(JvmToolLockfileSentinel):
     resolve_name = ScroogeSubsystem.options_scope
 
 
-@rule(desc="Generate Scala from Thrift", level=LogLevel.DEBUG)
-async def generate_scala_from_thrift_via_scrooge(
-    request: GenerateScalaFromThriftRequest,
+@rule
+async def generate_scrooge_thrift_sources(
+    request: GenerateScroogeThriftSourcesRequest,
     scrooge: ScroogeSubsystem,
     jdk_setup: JdkSetup,
     bash: BashBinary,
-) -> GeneratedSources:
+) -> GeneratedScroogeThriftSources:
     output_dir = "_generated_files"
     toolcp_relpath = "__toolcp"
 
-    tool_classpath, transitive_targets, empty_output_dir_digest = await MultiGet(
+    tool_classpath, transitive_targets, empty_output_dir_digest, wrapped_target = await MultiGet(
         Get(
             MaterializedClasspath,
             MaterializedClasspathRequest(
                 lockfiles=(scrooge.resolved_lockfile(),),
             ),
         ),
-        Get(TransitiveTargets, TransitiveTargetsRequest([request.protocol_target.address])),
+        Get(TransitiveTargets, TransitiveTargetsRequest([request.thrift_source_field.address])),
         Get(Digest, CreateDigest([Directory(output_dir)])),
+        Get(WrappedTarget, Address, request.thrift_source_field.address),
     )
 
     transitive_sources, target_sources = await MultiGet(
@@ -80,7 +68,7 @@ async def generate_scala_from_thrift_via_scrooge(
                 if tgt.has_field(ThriftSourceField)
             ),
         ),
-        Get(SourceFiles, SourceFilesRequest([request.protocol_target[ThriftSourceField]])),
+        Get(SourceFiles, SourceFilesRequest([request.thrift_source_field])),
     )
 
     sources_roots = await Get(
@@ -106,7 +94,7 @@ async def generate_scala_from_thrift_via_scrooge(
         maybe_include_paths.extend(["-i", path])
 
     maybe_finagle_option = []
-    if request.protocol_target[ScroogeFinagleBoolField].value:
+    if wrapped_target.target[ScroogeFinagleBoolField].value:
         maybe_finagle_option = ["--finagle"]
 
     immutable_input_digests = {
@@ -124,44 +112,23 @@ async def generate_scala_from_thrift_via_scrooge(
                 "--dest",
                 output_dir,
                 "--language",
-                "scala",
+                request.lang_id,
                 *maybe_finagle_option,
                 *target_sources.snapshot.files,
             ],
             input_digest=input_digest,
             immutable_input_digests=immutable_input_digests,
             use_nailgun=immutable_input_digests,
-            description=f"Generating Scala sources from {request.protocol_target.address}.",
+            description=f"Generating {request.lang_name} sources from {request.thrift_source_field.address}.",
             level=LogLevel.DEBUG,
             output_directories=(output_dir,),
             env=jdk_setup.env,
             append_only_caches=jdk_setup.append_only_caches,
         ),
     )
 
-    normalized_digest, source_root = await MultiGet(
-        Get(Digest, RemovePrefix(result.output_digest, output_dir)),
-        Get(SourceRoot, SourceRootRequest, SourceRootRequest.for_target(request.protocol_target)),
-    )
-
-    source_root_restored = (
-        await Get(Snapshot, AddPrefix(normalized_digest, source_root.path))
-        if source_root.path != "."
-        else await Get(Snapshot, Digest, normalized_digest)
-    )
-    return GeneratedSources(source_root_restored)
-
-
-class InjectScroogeDependencies(InjectDependenciesRequest):
-    inject_for = ThriftDependenciesField
-
-
-@rule
-async def inject_scrooge_dependencies(
-    _: InjectScroogeDependencies, scrooge: ScroogeSubsystem
-) -> InjectedDependencies:
-    addresses = await Get(Addresses, UnparsedAddressInputs, scrooge.runtime_dependencies)
-    return InjectedDependencies(addresses)
+    output_snapshot = await Get(Snapshot, RemovePrefix(result.output_digest, output_dir))
+    return GeneratedScroogeThriftSources(output_snapshot)
 
 
 @rule
@@ -176,7 +143,5 @@ def rules():
     return [
         *collect_rules(),
         *additional_fields.rules(),
-        UnionRule(GenerateSourcesRequest, GenerateScalaFromThriftRequest),
-        UnionRule(InjectDependenciesRequest, InjectScroogeDependencies),
         UnionRule(JvmToolLockfileSentinel, ScroogeToolLockfileSentinel),
     ]

src/python/pants/backend/codegen/thrift/scrooge/scala/BUILD

@@ -0,0 +1,6 @@
+# Copyright 2021 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+python_sources()
+
+python_tests(name="tests")

src/python/pants/backend/codegen/thrift/scrooge/scala/rules.py

@@ -0,0 +1,72 @@
+# Copyright 2021 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+from pants.backend.codegen.thrift.scrooge.rules import (
+    GeneratedScroogeThriftSources,
+    GenerateScroogeThriftSourcesRequest,
+)
+from pants.backend.codegen.thrift.scrooge.scala.subsystem import ScroogeScalaSubsystem
+from pants.backend.codegen.thrift.target_types import ThriftDependenciesField, ThriftSourceField
+from pants.backend.java.target_types import JavaSourceField
+from pants.engine.addresses import Addresses, UnparsedAddressInputs
+from pants.engine.fs import AddPrefix, Digest, Snapshot
+from pants.engine.internals.selectors import Get
+from pants.engine.rules import collect_rules, rule
+from pants.engine.target import (
+    GeneratedSources,
+    GenerateSourcesRequest,
+    InjectDependenciesRequest,
+    InjectedDependencies,
+)
+from pants.engine.unions import UnionRule
+from pants.source.source_root import SourceRoot, SourceRootRequest
+from pants.util.logging import LogLevel
+
+
+class GenerateScalaFromThriftRequest(GenerateSourcesRequest):
+    input = ThriftSourceField
+    output = JavaSourceField
+
+
+class InjectScroogeScalaDependencies(InjectDependenciesRequest):
+    inject_for = ThriftDependenciesField
+
+
+@rule(desc="Generate Scala from Thrift with Scrooge", level=LogLevel.DEBUG)
+async def generate_scala_from_thrift_with_scrooge(
+    request: GenerateScalaFromThriftRequest,
+) -> GeneratedSources:
+    result = await Get(
+        GeneratedScroogeThriftSources,
+        GenerateScroogeThriftSourcesRequest(
+            thrift_source_field=request.protocol_target[ThriftSourceField],
+            lang_id="scala",
+            lang_name="Scala",
+        ),
+    )
+
+    source_root = await Get(
+        SourceRoot, SourceRootRequest, SourceRootRequest.for_target(request.protocol_target)
+    )
+
+    source_root_restored = (
+        await Get(Snapshot, AddPrefix(result.snapshot.digest, source_root.path))
+        if source_root.path != "."
+        else await Get(Snapshot, Digest, result.snapshot.digest)
+    )
+    return GeneratedSources(source_root_restored)
+
+
+@rule
+async def inject_scrooge_scala_dependencies(
+    _: InjectScroogeScalaDependencies, scrooge: ScroogeScalaSubsystem
+) -> InjectedDependencies:
+    addresses = await Get(Addresses, UnparsedAddressInputs, scrooge.runtime_dependencies)
+    return InjectedDependencies(addresses)
+
+
+def rules():
+    return (
+        *collect_rules(),
+        UnionRule(GenerateSourcesRequest, GenerateScalaFromThriftRequest),
+        UnionRule(InjectDependenciesRequest, InjectScroogeScalaDependencies),
+    )

src/python/pants/backend/codegen/thrift/scrooge/rules_integration_test.py → src/python/pants/backend/codegen/thrift/scrooge/scala/rules_integration_test.py

@@ -7,8 +7,9 @@
 import pytest
 
 from pants.backend.codegen.thrift.rules import rules as thrift_rules
-from pants.backend.codegen.thrift.scrooge.rules import GenerateScalaFromThriftRequest
 from pants.backend.codegen.thrift.scrooge.rules import rules as scrooge_rules
+from pants.backend.codegen.thrift.scrooge.scala.rules import GenerateScalaFromThriftRequest
+from pants.backend.codegen.thrift.scrooge.scala.rules import rules as scrooge_scala_rules
 from pants.backend.codegen.thrift.target_types import (
     ThriftSourceField,
     ThriftSourcesGeneratorTarget,
@@ -35,6 +36,7 @@ def rule_runner() -> RuleRunner:
         rules=[
             *thrift_rules(),
             *scrooge_rules(),
+            *scrooge_scala_rules(),
             *config_files.rules(),
             *classpath.rules(),
             *coursier_fetch_rules(),

src/python/pants/backend/codegen/thrift/scrooge/scala/subsystem.py

@@ -0,0 +1,34 @@
+# Copyright 2021 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+from __future__ import annotations
+
+from pants.backend.codegen.thrift.scrooge.subsystem import ScroogeSubsystem
+from pants.engine.addresses import UnparsedAddressInputs
+from pants.option.custom_types import target_option
+from pants.option.subsystem import Subsystem
+
+
+class ScroogeScalaSubsystem(Subsystem):
+    options_scope = "scrooge-scala"
+    help = "Scala-specific options for the Scrooge Thrift IDL compiler (https://twitter.github.io/scrooge/)."
+
+    @classmethod
+    def register_options(cls, register):
+        super().register_options(register)
+        register(
+            "--runtime-dependencies",
+            type=list,
+            member_type=target_option,
+            help=(
+                "A list of addresses to `jvm_artifact` targets for the runtime "
+                "dependencies needed for generated Scala code to work. For example, "
+                "`['3rdparty/jvm:scrooge-runtime']`. These dependencies will "
+                "be automatically added to every `thrift_source` target. At the very least, "
+                "this option must be set to a `jvm_artifact` for the "
+                f"`com.twitter:scrooge-runtime_SCALAVER:{ScroogeSubsystem.default_version}` runtime library."
+            ),
+        )
+
+    @property
+    def runtime_dependencies(self) -> UnparsedAddressInputs:
+        return UnparsedAddressInputs(self.options.runtime_dependencies, owning_address=None)

src/python/pants/backend/codegen/thrift/scrooge/subsystem.py

@@ -2,9 +2,7 @@
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 from __future__ import annotations
 
-from pants.engine.addresses import UnparsedAddressInputs
 from pants.jvm.resolve.jvm_tool import JvmToolBase
-from pants.option.custom_types import target_option
 from pants.util.docutil import git_url
 
 
@@ -21,24 +19,3 @@ class ScroogeSubsystem(JvmToolBase):
     default_lockfile_url = git_url(
         "src/python/pants/backend/codegen/thrift/scrooge/scrooge.default.lockfile.txt"
     )
-
-    @classmethod
-    def register_options(cls, register):
-        super().register_options(register)
-        register(
-            "--runtime-dependencies",
-            type=list,
-            member_type=target_option,
-            help=(
-                "A list of addresses to `jvm_artifact` targets for the runtime "
-                "dependencies needed for generated Scala code to work. For example, "
-                "`['3rdparty/jvm:scrooge-runtime']`. These dependencies will "
-                "be automatically added to every `thrift_source` target. At the very least, "
-                "this option must be set to a `jvm_artifact` for the "
-                f"`com.twitter:scrooge-runtime_SCALAVER:{cls.default_version}` runtime library."
-            ),
-        )
-
-    @property
-    def runtime_dependencies(self) -> UnparsedAddressInputs:
-        return UnparsedAddressInputs(self.options.runtime_dependencies, owning_address=None)

src/python/pants/backend/experimental/codegen/thrift/scrooge/register.py → src/python/pants/backend/experimental/codegen/thrift/scrooge/scala/register.py

@@ -2,6 +2,7 @@
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 from pants.backend.codegen.thrift.rules import rules as thrift_rules
 from pants.backend.codegen.thrift.scrooge.rules import rules as scrooge_rules
+from pants.backend.codegen.thrift.scrooge.scala.rules import rules as scrooge_scala_rules
 from pants.backend.codegen.thrift.target_types import (
     ThriftSourcesGeneratorTarget,
     ThriftSourceTarget,
@@ -25,6 +26,7 @@ def rules():
     return [
         *thrift_rules(),
         *scrooge_rules(),
+        *scrooge_scala_rules(),
         # Re-export rules necessary to avoid rule graph errors.
         *config_files.rules(),
         *classpath.rules(),

src/python/pants/init/BUILD

@@ -12,7 +12,7 @@ target(
         "src/python/pants/backend/experimental/codegen/protobuf/scala",
         "src/python/pants/backend/experimental/codegen/thrift/apache/java",
         "src/python/pants/backend/experimental/codegen/thrift/apache/python",
-        "src/python/pants/backend/experimental/codegen/thrift/scrooge",
+        "src/python/pants/backend/experimental/codegen/thrift/scrooge/scala",
         "src/python/pants/backend/experimental/debian",
         "src/python/pants/backend/experimental/docker",
         "src/python/pants/backend/experimental/docker/lint/hadolint",