utils.py

import base64
import io
import textwrap
import uuid
from pathlib import Path
from typing import Any, override

from ai.backend.client.output.types import FieldSet, FieldSpec


def dedent(text: str) -> str:
    return textwrap.dedent(text).strip()


def create_connection_field(
    field_name: str, node_fields: FieldSet | tuple[FieldSpec, ...] | list[FieldSpec]
) -> FieldSpec:
    """
    Creates a GraphQL connection field specification with standard 'edges' and 'node' structure.
    Useful for building paginated queries where each edge contains a node with specified fields.
    """
    if isinstance(node_fields, (tuple, list)):
        node_fields = FieldSet(node_fields)

    return FieldSpec(
        field_name,
        subfields=FieldSet([
            FieldSpec("edges", subfields=FieldSet([FieldSpec("node", subfields=node_fields)]))
        ]),
    )


def flatten_connection(connection_data: dict[str, Any]) -> list[dict[str, Any]]:
    """
    Flattens a GraphQL Connection structure into a list of node dictionaries.
    Args:
        connection_data (dict): The GraphQL Connection data containing 'edges'.
    Returns:
        list[dict]: A list of node dictionaries extracted from the connection.
    """
    if connection_data is None or "edges" not in connection_data:
        return []
    return [edge["node"] for edge in connection_data["edges"]]


def flatten_connections_in_data(data: dict[str, Any] | None) -> dict[str, Any]:
    """
    Flattens all connection fields in a nested dictionary.
    If a value is a dictionary containing an 'edges' key, it is flattened using flatten_connection().
    Returns a new dictionary with all connections flattened.
    """
    result: dict[str, Any] = {}

    if data is None:
        return result

    for key, value in data.items():
        if key.endswith("_nodes") and isinstance(value, dict) and "edges" in value:
            result[key] = flatten_connection(value)
            continue
        result[key] = value

    return result


class ProgressReportingReader(io.BufferedReader):
    def __init__(self, file_path: str, *, tqdm_instance: Any = None) -> None:
        file_path_obj = Path(file_path)
        super().__init__(file_path_obj.open("rb"))
        self._filename = file_path_obj.name
        if tqdm_instance is None:
            from tqdm import tqdm

            self._owns_tqdm = True
            self.tqdm = tqdm(
                unit="bytes",
                unit_scale=True,
                total=file_path_obj.stat().st_size,
            )
        else:
            self._owns_tqdm = False
            self.tqdm = tqdm_instance

    @override
    def __enter__(self) -> "ProgressReportingReader":
        return self

    @override
    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        exc_traceback: Any,
    ) -> None:
        if self._owns_tqdm:
            self.tqdm.close()
        self.close()

    @override
    def read(self, *args: Any, **kwargs: Any) -> bytes:
        chunk = super().read(*args, **kwargs)
        self.tqdm.set_postfix(file=self._filename, refresh=False)
        self.tqdm.update(len(chunk))
        return chunk

    @override
    def read1(self, *args: Any, **kwargs: Any) -> bytes:
        chunk = super().read1(*args, **kwargs)
        self.tqdm.set_postfix(file=self._filename, refresh=False)
        self.tqdm.update(len(chunk))
        return chunk

    @override
    def readinto(self, *args: Any, **kwargs: Any) -> int:
        count = super().readinto(*args, **kwargs)
        self.tqdm.set_postfix(file=self._filename, refresh=False)
        self.tqdm.update(count)
        return count

    @override
    def readinto1(self, *args: Any, **kwargs: Any) -> int:
        count = super().readinto1(*args, **kwargs)
        self.tqdm.set_postfix(file=self._filename, refresh=False)
        self.tqdm.update(count)
        return count


def to_global_id(node_name: str, id: uuid.UUID) -> str:
    """
    Used to generate a global ID for a node in the GraphQL Relay specification.
    Encode the node name and id into a global ID using Base64 encoding.
    """
    return base64.b64encode(f"{node_name}:{id!s}".encode()).decode("ascii")