Add python_resolve field to protobuf_source and thrift_source t…

…o support multiple resolves with codegen (Cherry-pick of #14693) (#14698)

Closes #14484.

## Problem

Generated code needs associated runtime libraries, e.g. `protobuf`. Those must come from a `python_requirement` which by definition have a single `resolve`.

Before this PR, it was impossible to have >1 `python_requirement` target used for each runtime library and every `protobuf_source` target would depend on the same `python_requirement`. Practically, this means codegen could only work with code belonging to a single resolve.

## Solution

Add a `python_resolve` field to codegen targets so that they can indicate which `python_requirement` they should be using. Our new dependency inference from #14691 wires this up properly.

If you want the same generated code to work with multiple resolves, you will need to use `parametrize`. That gets tricky when you consider a codegen target generating for multiple languages, as explained at the bottom of #14484, but there is a decent workaround via configurations. That will be tackled in a followup.

### FYI: lazy validation of runtime library

#14691 asserts that there is exactly 1 runtime library in the project; now we assert that for each resolve. 

Key nuance: this check is lazy. If you have 2 resolves, but your codegen targets only use 1 of the 2, then we will never end up checking that there is a runtime library defined for the 2nd resolve. That is important so that, for example, if you have a `pants-plugins` resolve we don't force you to add `protobuf` unnecessarily to it.

[ci skip-rust]
[ci skip-build-wheels]

src/python/pants/backend/codegen/protobuf/python/additional_fields.py

@@ -5,14 +5,18 @@
     ProtobufSourcesGeneratorTarget,
     ProtobufSourceTarget,
 )
-from pants.backend.python.target_types import InterpreterConstraintsField
+from pants.backend.python.target_types import InterpreterConstraintsField, PythonResolveField
 from pants.engine.target import StringField
 
 
-class ProtobufPythonInterpreterConstraints(InterpreterConstraintsField):
+class ProtobufPythonInterpreterConstraintsField(InterpreterConstraintsField):
     alias = "python_interpreter_constraints"
 
 
+class ProtobufPythonResolveField(PythonResolveField):
+    alias = "python_resolve"
+
+
 class PythonSourceRootField(StringField):
     alias = "python_source_root"
     help = (
@@ -23,8 +27,12 @@ class PythonSourceRootField(StringField):
 
 def rules():
     return [
-        ProtobufSourceTarget.register_plugin_field(ProtobufPythonInterpreterConstraints),
-        ProtobufSourcesGeneratorTarget.register_plugin_field(ProtobufPythonInterpreterConstraints),
+        ProtobufSourceTarget.register_plugin_field(ProtobufPythonInterpreterConstraintsField),
+        ProtobufSourcesGeneratorTarget.register_plugin_field(
+            ProtobufPythonInterpreterConstraintsField
+        ),
+        ProtobufSourceTarget.register_plugin_field(ProtobufPythonResolveField),
+        ProtobufSourcesGeneratorTarget.register_plugin_field(ProtobufPythonResolveField),
         ProtobufSourceTarget.register_plugin_field(PythonSourceRootField),
         ProtobufSourcesGeneratorTarget.register_plugin_field(PythonSourceRootField),
     ]

src/python/pants/backend/codegen/protobuf/python/python_protobuf_subsystem.py

@@ -11,6 +11,8 @@
 from pants.backend.python.goals import lockfile
 from pants.backend.python.goals.lockfile import GeneratePythonLockfile
 from pants.backend.python.subsystems.python_tool_base import PythonToolRequirementsBase
+from pants.backend.python.subsystems.setup import PythonSetup
+from pants.backend.python.target_types import PythonResolveField
 from pants.core.goals.generate_lockfiles import GenerateToolLockfileSentinel
 from pants.engine.addresses import Address
 from pants.engine.rules import Get, collect_rules, rule
@@ -42,6 +44,10 @@ class PythonProtobufSubsystem(Subsystem):
             "`protobuf` module (usually from the `protobuf` requirement). If the `protobuf_source` "
             "target sets `grpc=True`, will also add a dependency on the `python_requirement` "
             "target exposing the `grpcio` module.\n\n"
+            "If `[python].enable_resolves` is set, Pants will only infer dependencies on "
+            "`python_requirement` targets that use the same resolve as the particular "
+            "`protobuf_source` / `protobuf_sources` target uses, which is set via its "
+            "`python_resolve` field.\n\n"
             "Unless this option is disabled, Pants will error if no relevant target is found or "
             "if more than one is found which causes ambiguity."
         ),
@@ -88,31 +94,39 @@ class InjectPythonProtobufDependencies(InjectDependenciesRequest):
 async def inject_dependencies(
     request: InjectPythonProtobufDependencies,
     python_protobuf: PythonProtobufSubsystem,
+    python_setup: PythonSetup,
     # TODO(#12946): Make this a lazy Get once possible.
     module_mapping: ThirdPartyPythonModuleMapping,
 ) -> InjectedDependencies:
     if not python_protobuf.infer_runtime_dependency:
         return InjectedDependencies()
 
+    wrapped_tgt = await Get(WrappedTarget, Address, request.dependencies_field.address)
+    tgt = wrapped_tgt.target
+    resolve = tgt.get(PythonResolveField).normalized_value(python_setup)
+
     result = [
         find_python_runtime_library_or_raise_error(
             module_mapping,
             request.dependencies_field.address,
             "google.protobuf",
+            resolve=resolve,
+            resolves_enabled=python_setup.enable_resolves,
             recommended_requirement_name="protobuf",
             recommended_requirement_url="https://pypi.org/project/protobuf/",
             disable_inference_option=f"[{python_protobuf.options_scope}].infer_runtime_dependency",
         )
     ]
 
-    wrapped_tgt = await Get(WrappedTarget, Address, request.dependencies_field.address)
     if wrapped_tgt.target.get(ProtobufGrpcToggleField).value:
         result.append(
             find_python_runtime_library_or_raise_error(
                 module_mapping,
                 request.dependencies_field.address,
                 # Note that the library is called `grpcio`, but the module is `grpc`.
                 "grpc",
+                resolve=resolve,
+                resolves_enabled=python_setup.enable_resolves,
                 recommended_requirement_name="grpcio",
                 recommended_requirement_url="https://pypi.org/project/grpcio/",
                 disable_inference_option=f"[{python_protobuf.options_scope}].infer_runtime_dependency",

src/python/pants/backend/codegen/protobuf/python/python_protobuf_subsystem_test.py

@@ -33,6 +33,9 @@ def test_find_protobuf_python_requirement() -> None:
     rule_runner.write_files(
         {"codegen/dir/f.proto": "", "codegen/dir/BUILD": "protobuf_sources(grpc=True)"}
     )
+    rule_runner.set_options(
+        ["--python-resolves={'python-default': '', 'another': ''}", "--python-enable-resolves"]
+    )
     proto_tgt = rule_runner.get_target(Address("codegen/dir", relative_file_path="f.proto"))
     request = InjectPythonProtobufDependencies(proto_tgt[Dependencies])
 
@@ -49,7 +52,20 @@ def test_find_protobuf_python_requirement() -> None:
         [Address("proto1"), Address("grpc1")]
     )
 
-    # If multiple, error.
+    # Multiple is fine if from other resolve.
+    rule_runner.write_files(
+        {
+            "another_resolve/BUILD": (
+                "python_requirement(name='r1', requirements=['protobuf'], resolve='another')\n"
+                "python_requirement(name='r2', requirements=['grpc'], resolve='another')\n"
+            )
+        }
+    )
+    assert rule_runner.request(InjectedDependencies, [request]) == InjectedDependencies(
+        [Address("proto1"), Address("grpc1")]
+    )
+
+    # If multiple from the same resolve, error.
     rule_runner.write_files({"grpc2/BUILD": "python_requirement(requirements=['grpc'])"})
     with engine_error(
         AmbiguousPythonCodegenRuntimeLibrary, contains="['grpc1:grpc1', 'grpc2:grpc2']"

src/python/pants/backend/codegen/thrift/apache/python/additional_fields.py

@@ -0,0 +1,19 @@
+# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from pants.backend.codegen.thrift.target_types import (
+    ThriftSourcesGeneratorTarget,
+    ThriftSourceTarget,
+)
+from pants.backend.python.target_types import PythonResolveField
+
+
+class ThriftPythonResolveField(PythonResolveField):
+    alias = "python_resolve"
+
+
+def rules():
+    return [
+        ThriftSourceTarget.register_plugin_field(ThriftPythonResolveField),
+        ThriftSourcesGeneratorTarget.register_plugin_field(ThriftPythonResolveField),
+    ]

src/python/pants/backend/codegen/thrift/apache/python/register.py

@@ -1,7 +1,10 @@
 # Copyright 2021 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
-from pants.backend.codegen.thrift.apache.python import python_thrift_module_mapper
+from pants.backend.codegen.thrift.apache.python import (
+    additional_fields,
+    python_thrift_module_mapper,
+)
 from pants.backend.codegen.thrift.apache.python.rules import rules as apache_thrift_python_rules
 from pants.backend.codegen.thrift.apache.rules import rules as apache_thrift_rules
 from pants.backend.codegen.thrift.rules import rules as thrift_rules
@@ -19,6 +22,7 @@ def target_types():
 
 def rules():
     return [
+        *additional_fields.rules(),
         *thrift_rules(),
         *apache_thrift_rules(),
         *apache_thrift_python_rules(),

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

@@ -12,15 +12,17 @@
 from pants.backend.codegen.thrift.target_types import ThriftDependenciesField, ThriftSourceField
 from pants.backend.codegen.utils import find_python_runtime_library_or_raise_error
 from pants.backend.python.dependency_inference.module_mapper import ThirdPartyPythonModuleMapping
-from pants.backend.python.target_types import PythonSourceField
+from pants.backend.python.subsystems.setup import PythonSetup
+from pants.backend.python.target_types import PythonResolveField, PythonSourceField
+from pants.engine.addresses import Address
 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.rules import Get, collect_rules, rule
 from pants.engine.target import (
     GeneratedSources,
     GenerateSourcesRequest,
     InjectDependenciesRequest,
     InjectedDependencies,
+    WrappedTarget,
 )
 from pants.engine.unions import UnionRule
 from pants.source.source_root import SourceRoot, SourceRootRequest
@@ -69,16 +71,22 @@ class InjectApacheThriftPythonDependencies(InjectDependenciesRequest):
 async def find_apache_thrift_python_requirement(
     request: InjectApacheThriftPythonDependencies,
     thrift_python: ThriftPythonSubsystem,
+    python_setup: PythonSetup,
     # TODO(#12946): Make this a lazy Get once possible.
     module_mapping: ThirdPartyPythonModuleMapping,
 ) -> InjectedDependencies:
     if not thrift_python.infer_runtime_dependency:
         return InjectedDependencies()
 
+    wrapped_tgt = await Get(WrappedTarget, Address, request.dependencies_field.address)
+    resolve = wrapped_tgt.target.get(PythonResolveField).normalized_value(python_setup)
+
     addr = find_python_runtime_library_or_raise_error(
         module_mapping,
         request.dependencies_field.address,
         "thrift",
+        resolve=resolve,
+        resolves_enabled=python_setup.enable_resolves,
         recommended_requirement_name="thrift",
         recommended_requirement_url="https://pypi.org/project/thrift/",
         disable_inference_option=f"[{thrift_python.options_scope}].infer_runtime_dependency",

src/python/pants/backend/codegen/thrift/apache/python/rules_integration_test.py

@@ -180,6 +180,9 @@ def test_top_level_source_root(rule_runner: RuleRunner) -> None:
 
 def test_find_thrift_python_requirement(rule_runner: RuleRunner) -> None:
     rule_runner.write_files({"codegen/dir/f.thrift": "", "codegen/dir/BUILD": "thrift_sources()"})
+    rule_runner.set_options(
+        ["--python-resolves={'python-default': '', 'another': ''}", "--python-enable-resolves"]
+    )
     thrift_tgt = rule_runner.get_target(Address("codegen/dir", relative_file_path="f.thrift"))
     request = InjectApacheThriftPythonDependencies(thrift_tgt[Dependencies])
 
@@ -193,7 +196,15 @@ def test_find_thrift_python_requirement(rule_runner: RuleRunner) -> None:
         [Address("reqs1")]
     )
 
-    # If multiple, error.
+    # Multiple is fine if from other resolve.
+    rule_runner.write_files(
+        {"another_resolve/BUILD": "python_requirement(requirements=['thrift'], resolve='another')"}
+    )
+    assert rule_runner.request(InjectedDependencies, [request]) == InjectedDependencies(
+        [Address("reqs1")]
+    )
+
+    # If multiple from the same resolve, error.
     rule_runner.write_files({"reqs2/BUILD": "python_requirement(requirements=['thrift'])"})
     with engine_error(
         AmbiguousPythonCodegenRuntimeLibrary, contains="['reqs1:reqs1', 'reqs2:reqs2']"

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

@@ -26,6 +26,10 @@ class ThriftPythonSubsystem(Subsystem):
         help=(
             "If True, will add a dependency on a `python_requirement` target exposing the `thrift` "
             "module (usually from the `thrift` requirement).\n\n"
+            "If `[python].enable_resolves` is set, Pants will only infer dependencies on "
+            "`python_requirement` targets that use the same resolve as the particular "
+            "`thrift_source` / `thrift_source` target uses, which is set via its "
+            "`python_resolve` field.\n\n"
             "Unless this option is disabled, Pants will error if no relevant target is found or "
             "more than one is found which causes ambiguity."
         ),

src/python/pants/backend/codegen/utils.py

@@ -24,39 +24,64 @@ def find_python_runtime_library_or_raise_error(
     codegen_address: Address,
     runtime_library_module: str,
     *,
+    resolve: str,
+    resolves_enabled: bool,
     recommended_requirement_name: str,
     recommended_requirement_url: str,
     disable_inference_option: str,
 ) -> Address:
     addresses = [
         module_provider.addr
         for module_provider in module_mapping.providers_for_module(
-            runtime_library_module, resolve=None
+            runtime_library_module, resolve=resolve
         )
         if module_provider.typ == ModuleProviderType.IMPL
     ]
+    if len(addresses) == 1:
+        return addresses[0]
 
+    for_resolve_str = f" for the resolve '{resolve}'" if resolves_enabled else ""
     if not addresses:
+        resolve_note = (
+            (
+                "Note that because `[python].enable_resolves` is set, you must specifically have a "
+                f"`python_requirement` target that uses the same resolve '{resolve}' as the target "
+                f"{codegen_address}. Alternatively, update {codegen_address} to use a different "
+                "resolve.\n\n"
+            )
+            if resolves_enabled
+            else ""
+        )
         raise MissingPythonCodegenRuntimeLibrary(
             f"No `python_requirement` target was found with the module `{runtime_library_module}` "
-            f"in your project, so the Python code generated from the target {codegen_address} will "
-            f"not work properly. See {doc_url('python-third-party-dependencies')} for how to "
+            f"in your project{for_resolve_str}, so the Python code generated from the target "
+            f"{codegen_address} will not work properly. See "
+            f"{doc_url('python-third-party-dependencies')} for how to "
             "add a requirement, such as adding to requirements.txt. Usually you will want to use "
             f"the `{recommended_requirement_name}` project at {recommended_requirement_url}.\n\n"
+            f"{resolve_note}"
             f"To ignore this error, set `{disable_inference_option} = false` in `pants.toml`."
         )
 
-    if len(addresses) > 1:
-        raise AmbiguousPythonCodegenRuntimeLibrary(
-            "Multiple `python_requirement` targets were found with the module "
-            f"`{runtime_library_module}` in your project, so it is ambiguous which to use for the "
-            f"runtime library for the Python code generated from the the target {codegen_address}: "
-            f"{sorted(addr.spec for addr in addresses)}\n\n"
-            "To fix, remove one of these `python_requirement` targets so that there is no "
-            "ambiguity and Pants can infer a dependency. Alternatively, if you do want to have "
+    alternative_solution = (
+        (
+            f"Alternatively, change the resolve field for {codegen_address} to use a different "
+            "resolve from `[python].resolves`."
+        )
+        if resolves_enabled
+        else (
+            "Alternatively, if you do want to have "
             f"multiple conflicting versions of the `{runtime_library_module}` requirement, set "
             f"`{disable_inference_option} = false` in `pants.toml`. "
             f"Then manually add a dependency on the relevant `python_requirement` target to each "
             "target that directly depends on this generated code (e.g. `python_source` targets)."
         )
-    return addresses[0]
+    )
+    raise AmbiguousPythonCodegenRuntimeLibrary(
+        "Multiple `python_requirement` targets were found with the module "
+        f"`{runtime_library_module}` in your project{for_resolve_str}, so it is ambiguous which to "
+        f"use for the runtime library for the Python code generated from the the target "
+        f"{codegen_address}: {sorted(addr.spec for addr in addresses)}\n\n"
+        f"To fix, remove one of these `python_requirement` targets{for_resolve_str} so that "
+        f"there is no ambiguity and Pants can infer a dependency. {alternative_solution}"
+    )