pantsbuild / pants / 19015773527

from __future__ import annotations

import os

import re

from collections import defaultdict

from collections.abc import Mapping

from dataclasses import dataclass

from typing import DefaultDict

from pants.backend.codegen.protobuf.target_types import AllProtobufTargets, ProtobufSourceField

from pants.engine.addresses import Address

from pants.engine.fs import FileContent

from pants.engine.internals.graph import hydrate_sources

from pants.engine.intrinsics import get_digest_contents

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

from pants.engine.target import HydrateSourcesRequest

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 FirstPartyProtobufJvmMappingRequest:

    capitalize_base_name: bool

@rule

async def map_first_party_protobuf_jvm_targets_to_symbols(

    sources = await concurrently(

    all_contents = await concurrently(

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

    for tgt, contents in zip(all_protobuf_targets, all_contents):

        if not contents:

            continue

        if len(contents) > 1:

            raise AssertionError(

        resolve = tgt[JvmResolveField].normalized_value(jvm)

        namespace = _determine_namespace(

        namespace_mapping[(resolve, namespace)].add(tgt.address)

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

    for (resolve, namespace), addresses in namespace_mapping.items():

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

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

def _determine_namespace(file: FileContent, *, capitalize_base_name: bool) -> str:

    base_name, _, _ = os.path.basename(file.path).partition(".")

    base_name = base_name.capitalize() if capitalize_base_name else base_name

    package_definition = _parse_package_definition(file.content)

    return f"{package_definition}.{base_name}" if package_definition else base_name

_QUOTE_CHAR = r"(?:'|\")"

_JAVA_PACKAGE_OPTION_RE = re.compile(

_PACKAGE_RE = re.compile(r"^\s*package\s+(.+);")

def _parse_package_definition(content_raw: bytes) -> str | None:

    content = content_raw.decode()

    for line in content.splitlines():

        m = _JAVA_PACKAGE_OPTION_RE.match(line)

        if m:

            return m.group(1)

        m = _PACKAGE_RE.match(line)

        if m:

            return m.group(1)

    return None

def rules():

    return collect_rules()