diff --git a/src/python/pants/backend/python/dependency_inference/module_mapper.py b/src/python/pants/backend/python/dependency_inference/module_mapper.py index f80aadca6ce..d34801bd6fa 100644 --- a/src/python/pants/backend/python/dependency_inference/module_mapper.py +++ b/src/python/pants/backend/python/dependency_inference/module_mapper.py @@ -4,6 +4,7 @@ from __future__ import annotations import enum +import itertools import logging from collections import defaultdict from dataclasses import dataclass @@ -16,7 +17,9 @@ DEFAULT_MODULE_MAPPING, DEFAULT_TYPE_STUB_MODULE_MAPPING, ) +from pants.backend.python.subsystems.setup import PythonSetup from pants.backend.python.target_types import ( + PythonRequirementCompatibleResolvesField, PythonRequirementModulesField, PythonRequirementsField, PythonRequirementTypeStubModulesField, @@ -173,10 +176,21 @@ async def map_first_party_python_targets_to_modules( # Third party module mapping # ----------------------------------------------------------------------------------------------- +_ResolveName = str -class ThirdPartyPythonModuleMapping(FrozenDict[str, Tuple[ModuleProvider, ...]]): - def providers_for_module(self, module: str) -> tuple[ModuleProvider, ...]: - result = self.get(module, ()) + +class ThirdPartyPythonModuleMapping( + FrozenDict[_ResolveName, FrozenDict[str, Tuple[ModuleProvider, ...]]] +): + """A mapping of each resolve to the modules they contain and the addresses providing those + modules.""" + + def _providers_for_resolve(self, module: str, resolve: str) -> tuple[ModuleProvider, ...]: + mapping = self.get(resolve) + if not mapping: + return () + + result = mapping.get(module, ()) if result: return result @@ -185,25 +199,50 @@ def providers_for_module(self, module: str) -> tuple[ModuleProvider, ...]: if "." not in module: return () parent_module = module.rsplit(".", maxsplit=1)[0] - return self.providers_for_module(parent_module) + return self._providers_for_resolve(parent_module, resolve) + + def providers_for_module( + self, module: str, resolves: Iterable[str] | None + ) -> tuple[ModuleProvider, ...]: + """Find all providers for the module. + + If `resolves` is None, will not consider resolves, i.e. any `python_requirement` can be + consumed. Otherwise, providers can only come from `python_requirements` marked compatible + with those resolves. + """ + if resolves is None: + resolves = list(self.keys()) + return tuple( + itertools.chain.from_iterable( + self._providers_for_resolve(module, resolve) for resolve in resolves + ) + ) @rule(desc="Creating map of third party targets to Python modules", level=LogLevel.DEBUG) async def map_third_party_modules_to_addresses( all_python_tgts: AllPythonTargets, + python_setup: PythonSetup, ) -> ThirdPartyPythonModuleMapping: - modules_to_providers: DefaultDict[str, list[ModuleProvider]] = defaultdict(list) + resolves_to_modules_to_providers: dict[ + _ResolveName, DefaultDict[str, list[ModuleProvider]] + ] = {} for tgt in all_python_tgts.third_party: + tgt[PythonRequirementCompatibleResolvesField].validate(python_setup) + resolves = tgt[PythonRequirementCompatibleResolvesField].value_or_default(python_setup) def add_modules(modules: Iterable[str], *, type_stub: bool = False) -> None: - for module in modules: - modules_to_providers[module].append( - ModuleProvider( - tgt.address, - ModuleProviderType.TYPE_STUB if type_stub else ModuleProviderType.IMPL, + for resolve in resolves: + if resolve not in resolves_to_modules_to_providers: + resolves_to_modules_to_providers[resolve] = defaultdict(list) + for module in modules: + resolves_to_modules_to_providers[resolve][module].append( + ModuleProvider( + tgt.address, + ModuleProviderType.TYPE_STUB if type_stub else ModuleProviderType.IMPL, + ) ) - ) explicit_modules = tgt.get(PythonRequirementModulesField).value if explicit_modules: @@ -240,7 +279,13 @@ def add_modules(modules: Iterable[str], *, type_stub: bool = False) -> None: add_modules(DEFAULT_MODULE_MAPPING.get(proj_name, (fallback_value,))) return ThirdPartyPythonModuleMapping( - (k, tuple(sorted(v))) for k, v in sorted(modules_to_providers.items()) + ( + resolve, + FrozenDict( + (mod, tuple(sorted(providers))) for mod, providers in sorted(mapping.items()) + ), + ) + for resolve, mapping in sorted(resolves_to_modules_to_providers.items()) ) @@ -276,7 +321,7 @@ async def map_module_to_address( third_party_mapping: ThirdPartyPythonModuleMapping, ) -> PythonModuleOwners: providers = [ - *third_party_mapping.providers_for_module(module.module), + *third_party_mapping.providers_for_module(module.module, resolves=None), *first_party_mapping.providers_for_module(module.module), ] addresses = tuple(provider.addr for provider in providers) diff --git a/src/python/pants/backend/python/dependency_inference/module_mapper_test.py b/src/python/pants/backend/python/dependency_inference/module_mapper_test.py index 9033fa1ea99..e1d8bb13a22 100644 --- a/src/python/pants/backend/python/dependency_inference/module_mapper_test.py +++ b/src/python/pants/backend/python/dependency_inference/module_mapper_test.py @@ -34,6 +34,7 @@ from pants.core.util_rules import stripped_source_files from pants.engine.addresses import Address from pants.testutil.rule_runner import QueryRule, RuleRunner +from pants.util.frozendict import FrozenDict def test_default_module_mapping_is_normalized() -> None: @@ -124,16 +125,23 @@ def test_third_party_modules_mapping() -> None: ) mapping = ThirdPartyPythonModuleMapping( { - "colors": (colors_provider, colors_stubs_provider), - "pants": (pants_provider,), - "req.submodule": (submodule_provider,), - "pants.testutil": (pants_testutil_provider,), - "ambiguous": (colors_provider, pants_provider), + "default-resolve": FrozenDict( + { + "colors": (colors_provider, colors_stubs_provider), + "pants": (pants_provider,), + "req.submodule": (submodule_provider,), + "pants.testutil": (pants_testutil_provider,), + "two_resolves": (colors_provider,), + } + ), + "another-resolve": FrozenDict({"two_resolves": (pants_provider,)}), } ) - def assert_addresses(mod: str, expected: tuple[ModuleProvider, ...]) -> None: - assert mapping.providers_for_module(mod) == expected + def assert_addresses( + mod: str, expected: tuple[ModuleProvider, ...], *, resolves: list[str] | None = None + ) -> None: + assert mapping.providers_for_module(mod, resolves) == expected assert_addresses("colors", (colors_provider, colors_stubs_provider)) assert_addresses("colors.red", (colors_provider, colors_stubs_provider)) @@ -154,9 +162,20 @@ def assert_addresses(mod: str, expected: tuple[ModuleProvider, ...]) -> None: assert_addresses("unknown", ()) assert_addresses("unknown.pants", ()) - assert_addresses("ambiguous", (colors_provider, pants_provider)) - assert_addresses("ambiguous.foo", (colors_provider, pants_provider)) - assert_addresses("ambiguous.foo.bar", (colors_provider, pants_provider)) + assert_addresses("two_resolves", (colors_provider, pants_provider), resolves=None) + assert_addresses("two_resolves.foo", (colors_provider, pants_provider), resolves=None) + assert_addresses("two_resolves.foo.bar", (colors_provider, pants_provider), resolves=None) + assert_addresses("two_resolves", (colors_provider,), resolves=["default-resolve"]) + assert_addresses("two_resolves", (pants_provider,), resolves=["another-resolve"]) + assert_addresses( + "two_resolves", + ( + colors_provider, + pants_provider, + ), + resolves=["default-resolve", "another-resolve"], + ) + assert_addresses("two_resolves", (), resolves=[]) @pytest.fixture @@ -280,11 +299,13 @@ def req( *, modules: list[str] | None = None, stub_modules: list[str] | None = None, + resolves: list[str] | None = None, ) -> str: return ( f"python_requirement(name='{tgt_name}', requirements=['{req_str}'], " f"modules={modules or []}," - f"type_stub_modules={stub_modules or []})" + f"type_stub_modules={stub_modules or []}," + f"experimental_compatible_resolves={resolves or ['default']})" ) build_file = "\n\n".join( @@ -302,59 +323,89 @@ def req( req("typed-dep5", "typed-dep5-foo", stub_modules=["typed_dep5"]), # A 3rd-party dependency can have both a type stub and implementation. req("multiple_owners1", "multiple_owners==1"), - req("multiple_owners2", "multiple_owners==2"), - req("multiple_owners_types", "types-multiple_owners==1"), + req("multiple_owners2", "multiple_owners==2", resolves=["another"]), + req("multiple_owners_types", "types-multiple_owners==1", resolves=["another"]), # Only assume it's a type stubs dep if we are certain it's not an implementation. req("looks_like_stubs", "looks-like-stubs-types", modules=["looks_like_stubs"]), ] ) rule_runner.write_files({"BUILD": build_file}) + rule_runner.set_options(["--python-experimental-resolves={'default': '', 'another': ''}"]) result = rule_runner.request(ThirdPartyPythonModuleMapping, []) assert result == ThirdPartyPythonModuleMapping( { - "file_dist": ( - ModuleProvider(Address("", target_name="file_dist"), ModuleProviderType.IMPL), - ), - "looks_like_stubs": ( - ModuleProvider( - Address("", target_name="looks_like_stubs"), ModuleProviderType.IMPL - ), - ), - "mapped_module": ( - ModuleProvider(Address("", target_name="modules"), ModuleProviderType.IMPL), - ), - "multiple_owners": ( - ModuleProvider( - Address("", target_name="multiple_owners1"), ModuleProviderType.IMPL - ), - ModuleProvider( - Address("", target_name="multiple_owners2"), ModuleProviderType.IMPL - ), - ModuleProvider( - Address("", target_name="multiple_owners_types"), ModuleProviderType.TYPE_STUB - ), - ), - "req1": (ModuleProvider(Address("", target_name="req1"), ModuleProviderType.IMPL),), - "typed_dep1": ( - ModuleProvider(Address("", target_name="typed-dep1"), ModuleProviderType.TYPE_STUB), - ), - "typed_dep2": ( - ModuleProvider(Address("", target_name="typed-dep2"), ModuleProviderType.TYPE_STUB), - ), - "typed_dep3": ( - ModuleProvider(Address("", target_name="typed-dep3"), ModuleProviderType.TYPE_STUB), - ), - "typed_dep4": ( - ModuleProvider(Address("", target_name="typed-dep4"), ModuleProviderType.TYPE_STUB), - ), - "typed_dep5": ( - ModuleProvider(Address("", target_name="typed-dep5"), ModuleProviderType.TYPE_STUB), - ), - "un_normalized_project": ( - ModuleProvider(Address("", target_name="un_normalized"), ModuleProviderType.IMPL), + "another": FrozenDict( + { + "multiple_owners": ( + ModuleProvider( + Address("", target_name="multiple_owners2"), ModuleProviderType.IMPL + ), + ModuleProvider( + Address("", target_name="multiple_owners_types"), + ModuleProviderType.TYPE_STUB, + ), + ), + } ), - "vcs_dist": ( - ModuleProvider(Address("", target_name="vcs_dist"), ModuleProviderType.IMPL), + "default": FrozenDict( + { + "file_dist": ( + ModuleProvider( + Address("", target_name="file_dist"), ModuleProviderType.IMPL + ), + ), + "looks_like_stubs": ( + ModuleProvider( + Address("", target_name="looks_like_stubs"), ModuleProviderType.IMPL + ), + ), + "mapped_module": ( + ModuleProvider(Address("", target_name="modules"), ModuleProviderType.IMPL), + ), + "multiple_owners": ( + ModuleProvider( + Address("", target_name="multiple_owners1"), ModuleProviderType.IMPL + ), + ), + "req1": ( + ModuleProvider(Address("", target_name="req1"), ModuleProviderType.IMPL), + ), + "typed_dep1": ( + ModuleProvider( + Address("", target_name="typed-dep1"), ModuleProviderType.TYPE_STUB + ), + ), + "typed_dep2": ( + ModuleProvider( + Address("", target_name="typed-dep2"), ModuleProviderType.TYPE_STUB + ), + ), + "typed_dep3": ( + ModuleProvider( + Address("", target_name="typed-dep3"), ModuleProviderType.TYPE_STUB + ), + ), + "typed_dep4": ( + ModuleProvider( + Address("", target_name="typed-dep4"), ModuleProviderType.TYPE_STUB + ), + ), + "typed_dep5": ( + ModuleProvider( + Address("", target_name="typed-dep5"), ModuleProviderType.TYPE_STUB + ), + ), + "un_normalized_project": ( + ModuleProvider( + Address("", target_name="un_normalized"), ModuleProviderType.IMPL + ), + ), + "vcs_dist": ( + ModuleProvider( + Address("", target_name="vcs_dist"), ModuleProviderType.IMPL + ), + ), + } ), } ) diff --git a/src/python/pants/backend/python/goals/lockfile.py b/src/python/pants/backend/python/goals/lockfile.py index 2a3910f8e25..5f00ad4e17f 100644 --- a/src/python/pants/backend/python/goals/lockfile.py +++ b/src/python/pants/backend/python/goals/lockfile.py @@ -24,7 +24,7 @@ from pants.backend.python.subsystems.setup import PythonSetup from pants.backend.python.target_types import ( EntryPoint, - PythonCompatibleResolvesField, + PythonRequirementCompatibleResolvesField, PythonRequirementsField, UnrecognizedResolveNamesError, ) @@ -273,10 +273,10 @@ async def setup_user_lockfile_requests( resolve_to_requirements_fields = defaultdict(set) for tgt in all_targets: - if not tgt.has_field(PythonCompatibleResolvesField): + if not tgt.has_field(PythonRequirementCompatibleResolvesField): continue - tgt[PythonCompatibleResolvesField].validate(python_setup) - for resolve in tgt[PythonCompatibleResolvesField].value_or_default(python_setup): + tgt[PythonRequirementCompatibleResolvesField].validate(python_setup) + for resolve in tgt[PythonRequirementCompatibleResolvesField].value_or_default(python_setup): resolve_to_requirements_fields[resolve].add(tgt[PythonRequirementsField]) # TODO: Figure out how to determine which interpreter constraints to use for each resolve... diff --git a/src/python/pants/backend/python/target_types.py b/src/python/pants/backend/python/target_types.py index a17aeb303df..509d0a7e285 100644 --- a/src/python/pants/backend/python/target_types.py +++ b/src/python/pants/backend/python/target_types.py @@ -1022,7 +1022,7 @@ class PythonRequirementTarget(Target): PythonRequirementsField, PythonRequirementModulesField, PythonRequirementTypeStubModulesField, - PythonCompatibleResolvesField, + PythonRequirementCompatibleResolvesField, ) help = ( "A Python requirement installable by pip.\n\n"