Skip to content

Commit

Permalink
Python dependency inference supports multiple resolves for first-part…
Browse files Browse the repository at this point in the history 
…y targets (Cherry-pick of #14486) (#14765)

If you have >1 target for the same first-party module, that is now okay with dependency inference as long as they each have a distinct resolve. This is key for `parametrize()` to work properly:

```python
python_source(
    name="foo",
    source="foo.py",
    resolve=parametrize("a", "b"),
)
```

Note how codegen targets should have a `python_resolve` field now, per #14693. That allows support for multiple runtime libraries like `thrift` and `protobuf` dependencies. Now Python will only infer deps on codegen targets using the correct resolve for its runtime library. 

[ci skip-rust]
[ci skip-build-wheels]
  • Loading branch information
@Eric-Arellano
Eric-Arellano authored Mar 12, 2022
1 parent d9cf5c3 commit d418792
Show file tree
Hide file tree
Showing 6 changed files with 322 additions and 163 deletions.
25 changes: 17 additions & 8 deletions src/python/pants/backend/codegen/protobuf/python/python_protobuf_module_mapper.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,10 @@
FirstPartyPythonMappingImplMarker,
ModuleProvider,
ModuleProviderType,
ResolveName,
)
from pants.backend.python.subsystems.setup import PythonSetup
from pants.backend.python.target_types import PythonResolveField
from pants.core.util_rules.stripped_source_files import StrippedFileName, StrippedFileNameRequest
from pants.engine.rules import Get, MultiGet, collect_rules, rule
from pants.engine.unions import UnionRule
Expand All @@ -33,29 +36,35 @@ class PythonProtobufMappingMarker(FirstPartyPythonMappingImplMarker):
@rule(desc="Creating map of Protobuf targets to generated Python modules", level=LogLevel.DEBUG)
async def map_protobuf_to_python_modules(
protobuf_targets: AllProtobufTargets,
python_setup: PythonSetup,
_: PythonProtobufMappingMarker,
) -> FirstPartyPythonMappingImpl:

stripped_file_per_target = await MultiGet(
Get(StrippedFileName, StrippedFileNameRequest(tgt[ProtobufSourceField].file_path))
for tgt in protobuf_targets
)

modules_to_providers: DefaultDict[str, list[ModuleProvider]] = defaultdict(list)
resolves_to_modules_to_providers: DefaultDict[
ResolveName, DefaultDict[str, list[ModuleProvider]]
] = defaultdict(lambda: defaultdict(list))
for tgt, stripped_file in zip(protobuf_targets, stripped_file_per_target):
resolve = tgt[PythonResolveField].normalized_value(python_setup)

# NB: We don't consider the MyPy plugin, which generates `_pb2.pyi`. The stubs end up
# sharing the same module as the implementation `_pb2.py`. Because both generated files
# come from the same original Protobuf target, we're covered.
modules_to_providers[proto_path_to_py_module(stripped_file.value, suffix="_pb2")].append(
module = proto_path_to_py_module(stripped_file.value, suffix="_pb2")
resolves_to_modules_to_providers[resolve][module].append(
ModuleProvider(tgt.address, ModuleProviderType.IMPL)
)
if tgt.get(ProtobufGrpcToggleField).value:
modules_to_providers[
proto_path_to_py_module(stripped_file.value, suffix="_pb2_grpc")
].append(ModuleProvider(tgt.address, ModuleProviderType.IMPL))
module = proto_path_to_py_module(stripped_file.value, suffix="_pb2_grpc")
resolves_to_modules_to_providers[resolve][module].append(
ModuleProvider(tgt.address, ModuleProviderType.IMPL)
)

return FirstPartyPythonMappingImpl(
(k, tuple(sorted(v))) for k, v in sorted(modules_to_providers.items())
)
return FirstPartyPythonMappingImpl.create(resolves_to_modules_to_providers)


def rules():
Expand Down
65 changes: 47 additions & 18 deletions src/python/pants/backend/codegen/protobuf/python/python_protobuf_module_mapper_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@

from __future__ import annotations

from textwrap import dedent

import pytest

from pants.backend.codegen.protobuf.python import additional_fields, python_protobuf_module_mapper
Expand Down Expand Up @@ -36,7 +38,13 @@ def rule_runner() -> RuleRunner:


def test_map_first_party_modules_to_addresses(rule_runner: RuleRunner) -> None:
rule_runner.set_options(["--source-root-patterns=['root1', 'root2', 'root3']"])
rule_runner.set_options(
[
"--source-root-patterns=['root1', 'root2', 'root3']",
"--python-enable-resolves",
"--python-resolves={'python-default': '', 'another-resolve': ''}",
]
)
rule_runner.write_files(
{
"root1/protos/f1.proto": "",
Expand All @@ -47,41 +55,62 @@ def test_map_first_party_modules_to_addresses(rule_runner: RuleRunner) -> None:
"root1/two_owners/BUILD": "protobuf_sources()",
"root2/two_owners/f.proto": "",
"root2/two_owners/BUILD": "protobuf_sources()",
# A file with grpc. This also uses the `python_source_root` mechanism, which should be
# irrelevant to the module mapping because we strip source roots.
"root1/tests/f.proto": "",
"root1/tests/BUILD": "protobuf_sources(grpc=True, python_source_root='root3')",
"root1/tests/BUILD": dedent(
"""\
protobuf_sources(
grpc=True,
# This should be irrelevant to the module mapping because we strip source roots.
python_source_root='root3',
python_resolve='another-resolve',
)
"""
),
}
)
result = rule_runner.request(FirstPartyPythonMappingImpl, [PythonProtobufMappingMarker()])

def providers(addresses: list[Address]) -> tuple[ModuleProvider, ...]:
return tuple(ModuleProvider(addr, ModuleProviderType.IMPL) for addr in addresses)

assert result == FirstPartyPythonMappingImpl(
assert result == FirstPartyPythonMappingImpl.create(
{
"protos.f1_pb2": providers([Address("root1/protos", relative_file_path="f1.proto")]),
"protos.f2_pb2": providers([Address("root1/protos", relative_file_path="f2.proto")]),
"tests.f_pb2": providers([Address("root1/tests", relative_file_path="f.proto")]),
"tests.f_pb2_grpc": providers([Address("root1/tests", relative_file_path="f.proto")]),
"two_owners.f_pb2": providers(
[
Address("root1/two_owners", relative_file_path="f.proto"),
Address("root2/two_owners", relative_file_path="f.proto"),
]
),
"python-default": {
"protos.f1_pb2": providers(
[Address("root1/protos", relative_file_path="f1.proto")]
),
"protos.f2_pb2": providers(
[Address("root1/protos", relative_file_path="f2.proto")]
),
"two_owners.f_pb2": providers(
[
Address("root1/two_owners", relative_file_path="f.proto"),
Address("root2/two_owners", relative_file_path="f.proto"),
]
),
},
"another-resolve": {
"tests.f_pb2": providers([Address("root1/tests", relative_file_path="f.proto")]),
"tests.f_pb2_grpc": providers(
[Address("root1/tests", relative_file_path="f.proto")]
),
},
}
)


def test_top_level_source_root(rule_runner: RuleRunner) -> None:
rule_runner.set_options(["--source-root-patterns=['/']"])
rule_runner.set_options(["--source-root-patterns=['/']", "--python-enable-resolves"])
rule_runner.write_files({"protos/f.proto": "", "protos/BUILD": "protobuf_sources()"})
result = rule_runner.request(FirstPartyPythonMappingImpl, [PythonProtobufMappingMarker()])

def providers(addresses: list[Address]) -> tuple[ModuleProvider, ...]:
return tuple(ModuleProvider(addr, ModuleProviderType.IMPL) for addr in addresses)

assert result == FirstPartyPythonMappingImpl(
{"protos.f_pb2": providers([Address("protos", relative_file_path="f.proto")])}
assert result == FirstPartyPythonMappingImpl.create(
{
"python-default": {
"protos.f_pb2": providers([Address("protos", relative_file_path="f.proto")])
}
}
)
19 changes: 13 additions & 6 deletions src/python/pants/backend/codegen/thrift/apache/python/python_thrift_module_mapper.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,10 @@
FirstPartyPythonMappingImplMarker,
ModuleProvider,
ModuleProviderType,
ResolveName,
)
from pants.backend.python.subsystems.setup import PythonSetup
from pants.backend.python.target_types import PythonResolveField
from pants.engine.rules import Get, MultiGet, collect_rules, rule
from pants.engine.unions import UnionRule
from pants.util.logging import LogLevel
Expand All @@ -33,22 +36,26 @@ class PythonThriftMappingMarker(FirstPartyPythonMappingImplMarker):
@rule(desc="Creating map of Thrift targets to generated Python modules", level=LogLevel.DEBUG)
async def map_thrift_to_python_modules(
thrift_targets: AllThriftTargets,
python_setup: PythonSetup,
_: PythonThriftMappingMarker,
) -> FirstPartyPythonMappingImpl:
parsed_files = await MultiGet(
Get(ParsedThrift, ParsedThriftRequest(tgt[ThriftSourceField])) for tgt in thrift_targets
)
modules_to_providers: DefaultDict[str, list[ModuleProvider]] = defaultdict(list)

resolves_to_modules_to_providers: DefaultDict[
ResolveName, DefaultDict[str, list[ModuleProvider]]
] = defaultdict(lambda: defaultdict(list))
for tgt, parsed in zip(thrift_targets, parsed_files):
resolve = tgt[PythonResolveField].normalized_value(python_setup)
provider = ModuleProvider(tgt.address, ModuleProviderType.IMPL)
m1, m2 = thrift_path_to_py_modules(
source_path=tgt[ThriftSourceField].file_path, namespace=parsed.namespaces.get("py")
)
modules_to_providers[m1].append(provider)
modules_to_providers[m2].append(provider)
return FirstPartyPythonMappingImpl(
(k, tuple(sorted(v))) for k, v in sorted(modules_to_providers.items())
)
resolves_to_modules_to_providers[resolve][m1].append(provider)
resolves_to_modules_to_providers[resolve][m2].append(provider)

return FirstPartyPythonMappingImpl.create(resolves_to_modules_to_providers)


def rules():
Expand Down
84 changes: 55 additions & 29 deletions src/python/pants/backend/codegen/thrift/apache/python/python_thrift_module_mappper_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,14 @@

from __future__ import annotations

from textwrap import dedent

import pytest

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.python_thrift_module_mapper import (
PythonThriftMappingMarker,
)
Expand All @@ -24,6 +29,7 @@
def rule_runner() -> RuleRunner:
return RuleRunner(
rules=[
*additional_fields.rules(),
*python_thrift_module_mapper.rules(),
*thrift_target_types_rules(),
QueryRule(FirstPartyPythonMappingImpl, [PythonThriftMappingMarker]),
Expand All @@ -33,12 +39,26 @@ def rule_runner() -> RuleRunner:


def test_map_first_party_modules_to_addresses(rule_runner: RuleRunner) -> None:
rule_runner.set_options(["--source-root-patterns=['root1', 'root2', 'root3']"])
rule_runner.set_options(
[
"--source-root-patterns=['root1', 'root2']",
"--python-enable-resolves",
"--python-resolves={'python-default': '', 'another-resolve': ''}",
]
)
rule_runner.write_files(
{
"root1/thrift/no_namespace.thrift": "",
"root1/thrift/namespace.thrift": "namespace py custom_namespace.module",
"root1/thrift/BUILD": "thrift_sources()",
"root1/thrift/BUILD": dedent(
"""\
thrift_sources(
overrides={
"no_namespace.thrift": {"python_resolve": "another-resolve"}
},
)
"""
),
# These files will result in the same module name.
"root1/foo/two_owners.thrift": "",
"root1/foo/BUILD": "thrift_sources()",
Expand All @@ -51,29 +71,35 @@ def test_map_first_party_modules_to_addresses(rule_runner: RuleRunner) -> None:
def providers(addresses: list[Address]) -> tuple[ModuleProvider, ...]:
return tuple(ModuleProvider(addr, ModuleProviderType.IMPL) for addr in addresses)

assert dict(result) == {
"no_namespace.constants": providers(
[Address("root1/thrift", relative_file_path="no_namespace.thrift")]
),
"no_namespace.ttypes": providers(
[Address("root1/thrift", relative_file_path="no_namespace.thrift")]
),
"custom_namespace.module.constants": providers(
[Address("root1/thrift", relative_file_path="namespace.thrift")]
),
"custom_namespace.module.ttypes": providers(
[Address("root1/thrift", relative_file_path="namespace.thrift")]
),
"two_owners.constants": providers(
[
Address("root1/foo", relative_file_path="two_owners.thrift"),
Address("root2/bar", relative_file_path="two_owners.thrift"),
]
),
"two_owners.ttypes": providers(
[
Address("root1/foo", relative_file_path="two_owners.thrift"),
Address("root2/bar", relative_file_path="two_owners.thrift"),
]
),
}
assert result == FirstPartyPythonMappingImpl.create(
{
"python-default": {
"custom_namespace.module.constants": providers(
[Address("root1/thrift", relative_file_path="namespace.thrift")]
),
"custom_namespace.module.ttypes": providers(
[Address("root1/thrift", relative_file_path="namespace.thrift")]
),
"two_owners.constants": providers(
[
Address("root1/foo", relative_file_path="two_owners.thrift"),
Address("root2/bar", relative_file_path="two_owners.thrift"),
]
),
"two_owners.ttypes": providers(
[
Address("root1/foo", relative_file_path="two_owners.thrift"),
Address("root2/bar", relative_file_path="two_owners.thrift"),
]
),
},
"another-resolve": {
"no_namespace.constants": providers(
[Address("root1/thrift", relative_file_path="no_namespace.thrift")]
),
"no_namespace.ttypes": providers(
[Address("root1/thrift", relative_file_path="no_namespace.thrift")]
),
},
},
)
Loading

0 comments on commit d418792

Please sign in to comment.