pantsbuild / pants / 19015773527

from __future__ import annotations

from abc import ABCMeta

from collections import defaultdict

from collections.abc import Callable, Iterable, Iterator, Mapping

from dataclasses import dataclass

from pathlib import PurePath

from typing import Any, Generic, TypeVar

from pants.engine.collection import Collection

from pants.util.frozendict import FrozenDict

@dataclass(unsafe_hash=True)

class YamlPath:

    _elements: tuple[str, ...]

    _absolute: bool

    def __init__(self, elements: Iterable[str], *, absolute: bool) -> None:

        object.__setattr__(self, "_elements", tuple(elements))

        object.__setattr__(self, "_absolute", absolute)

        if len(self._elements) == 0 and not self._absolute:

            raise ValueError("Relative YAML paths with no elements are not allowed.")

    @classmethod

    def parse(cls, path: str) -> YamlPath:

        is_absolute = path.startswith("/")

        return cls([elem for elem in path.split("/") if elem], absolute=is_absolute)

    @classmethod

    def root(cls) -> YamlPath:

        return cls([], absolute=True)

    @classmethod

    def index(cls, idx: int) -> YamlPath:

        return cls([str(idx)], absolute=False)

    @property

    def parent(self) -> YamlPath | None:

        if not self.is_root:

            return YamlPath(self._elements[:-1], absolute=self._absolute)

        return None

    @property

    def current(self) -> str:

        if self.is_root:

            return ""

        return self._elements[len(self._elements) - 1]

    @property

    def is_absolute(self) -> bool:

        return self._absolute

    @property

    def is_root(self) -> bool:

        return len(self._elements) == 0 and self._absolute

    @property

    def is_index(self) -> bool:

        try:

            int(self.current)

            return True

        except ValueError:

            return False

    def to_relative(self) -> YamlPath:

        if not self._absolute:

            return self

        return YamlPath(self._elements, absolute=False)

    def __truediv__(self, other: str | int | YamlPath) -> YamlPath:

        if isinstance(other, str):

            other_path = YamlPath.parse(other)

        elif isinstance(other, int):

            other_path = YamlPath.index(other)

            other_path = other

        if other_path._absolute:

            raise ValueError("Can not append an absolute path to another path.")

        return YamlPath(self._elements + other_path._elements, absolute=self._absolute)

    def __iter__(self):

        return iter(self._elements)

    def __str__(self) -> str:

        path = "/".join(self._elements)

        if self._absolute:

            path = f"/{path}"

        return path

@dataclass(frozen=True)

class YamlElement(metaclass=ABCMeta):

    element_path: YamlPath

T = TypeVar("T")

R = TypeVar("R")

class MutableYamlIndex(Generic[T]):

    _data: dict[PurePath, dict[int, dict[YamlPath, T]]]

    def __init__(self) -> None:

        self._data = defaultdict(dict)

    def insert(

        doc_index = self._data[file_path].get(document_index, {})

        if not doc_index:

            self._data[file_path][document_index] = doc_index

        doc_index[yaml_path] = item

    def frozen(self) -> FrozenYamlIndex[T]:

        return FrozenYamlIndex.create(self)

@dataclass(frozen=True)

class _YamlDocumentIndexNode(Generic[T]):

    paths: FrozenDict[YamlPath, T]

    @classmethod

    def empty(cls: type[_YamlDocumentIndexNode[T]]) -> _YamlDocumentIndexNode[T]:

        return cls(paths=FrozenDict())

    def to_json_dict(self) -> dict[str, dict[str, str]]:

        items_dict: dict[str, str] = {}

        for path, item in self.paths.items():

            items_dict[str(path)] = str(item)

        return {"paths": items_dict}

@dataclass(frozen=True)

class FrozenYamlIndex(Generic[T]):

    _data: FrozenDict[PurePath, Collection[_YamlDocumentIndexNode[T]]]

    @classmethod

    def create(cls, other: MutableYamlIndex[T]) -> FrozenYamlIndex[T]:

        data: dict[PurePath, Collection[_YamlDocumentIndexNode[T]]] = {}

        for file_path, doc_index in other._data.items():

            max_index = max(doc_index.keys())

            doc_list: list[_YamlDocumentIndexNode[T]] = [_YamlDocumentIndexNode.empty()] * (

            for idx, item_map in doc_index.items():

                doc_list[idx] = _YamlDocumentIndexNode(paths=FrozenDict(item_map))

            data[file_path] = Collection(doc_list)

        return FrozenYamlIndex(_data=FrozenDict(data))

    @classmethod

    def empty(cls: type[FrozenYamlIndex[T]]) -> FrozenYamlIndex[T]:

        return FrozenYamlIndex[T](_data=FrozenDict())

    def transform_values(self, func: Callable[[T], R | None]) -> FrozenYamlIndex[R]:

        mutable_index: MutableYamlIndex[R] = MutableYamlIndex()

        for file_path, doc_index, yaml_path, item in self:

            new_item = func(item)

            if new_item is not None:

                mutable_index.insert(

        return mutable_index.frozen()

    def values(self) -> Iterator[T]:

        for _, _, _, item in self:

            yield item

    def to_json_dict(self) -> dict[str, Any]:

        result = {}

        for file_path, documents in self._data.items():

            result[str(file_path)] = [doc_idx.to_json_dict() for doc_idx in documents]

        return result

    def __iter__(self):

        for file_path, doc_indexes in self._data.items():

            for idx, doc_index in enumerate(doc_indexes):

                for yaml_path, item in doc_index.paths.items():

                    yield file_path, idx, yaml_path, item

def _to_snake_case(str: str) -> str:

    base_string = str.replace("-", "_")

    result = ""

    idx = 0

    for c in base_string:

        char_to_add = c

        if char_to_add.isupper():

            char_to_add = c.lower()

            if idx > 0:

                result += "_"

        result += char_to_add

        idx += 1

    return result

def snake_case_attr_dict(d: Mapping[str, Any]) -> dict[str, Any]:

    return {_to_snake_case(name): value for name, value in d.items()}