pantsbuild / pants / 19015773527

from __future__ import annotations

from collections import defaultdict

from collections.abc import Mapping

from dataclasses import dataclass

from typing import DefaultDict

from pants.backend.codegen.thrift.target_types import AllThriftTargets, ThriftSourceField

from pants.backend.codegen.thrift.thrift_parser import ParsedThriftRequest, parse_thrift_file

from pants.engine.addresses import Address

from pants.engine.rules import collect_rules, concurrently, rule

from pants.jvm.dependency_inference.artifact_mapper import MutableTrieNode

from pants.jvm.dependency_inference.symbol_mapper import SymbolMap

from pants.jvm.subsystems import JvmSubsystem

from pants.jvm.target_types import JvmResolveField

from pants.util.ordered_set import OrderedSet

_ResolveName = str

@dataclass(frozen=True)

class FirstPartyThriftJvmMappingRequest:

    extra_lang_ids: tuple[str, ...] = ()

    extra_namespace_directives: tuple[str, ...] = ()

@rule

async def map_first_party_thrift_targets_to_jvm_symbols(

    jvm_thrift_targets = [tgt for tgt in thrift_targets if tgt.has_field(JvmResolveField)]

    parsed_thrift_sources = await concurrently(

    package_symbols: DefaultDict[tuple[_ResolveName, str], OrderedSet[Address]] = defaultdict(

    lang_ids = {"java", *request.extra_lang_ids}

    for target, parsed_thrift in zip(jvm_thrift_targets, parsed_thrift_sources):

        for lang_id in lang_ids:

            package_name = parsed_thrift.namespaces.get(lang_id)

            if not package_name:

                continue

            resolve = target[JvmResolveField].normalized_value(jvm)

            package_symbols[(resolve, package_name)].add(target.address)

    mapping: Mapping[str, MutableTrieNode] = defaultdict(MutableTrieNode)

    for (resolve, package_name), addresses in package_symbols.items():

        mapping[resolve].insert(package_name, addresses, first_party=True, recursive=True)

    return SymbolMap((resolve, node.frozen()) for resolve, node in mapping.items())

def rules():

    return collect_rules()