Source code for tamr_client.attributes.attribute_type

"""
See https://docs.tamr.com/reference#attribute-types
"""
from dataclasses import dataclass
import logging
from typing import ClassVar, Tuple, Union

import tamr_client as tc
from tamr_client.types import JsonDict

logger = logging.getLogger(__name__)

# primitive types
#################


@dataclass(frozen=True)
class Boolean:
    _tag: ClassVar[str] = "BOOLEAN"


@dataclass(frozen=True)
class Double:
    _tag: ClassVar[str] = "DOUBLE"


@dataclass(frozen=True)
class Int:
    _tag: ClassVar[str] = "INT"


@dataclass(frozen=True)
class Long:
    _tag: ClassVar[str] = "LONG"


@dataclass(frozen=True)
class String:
    _tag: ClassVar[str] = "STRING"


PrimitiveType = Union[Boolean, Double, Int, Long, String]

# complex types
###############


@dataclass(frozen=True)
class Array:
    """See https://docs.tamr.com/reference#attribute-types"""

    # NOTE(pcattori) sphinx_autodoc_typehints cannot handle recursive reference
    # docstring written manually
    _tag: ClassVar[str] = "ARRAY"
    inner_type: "AttributeType"


@dataclass(frozen=True)
class Map:
    """See https://docs.tamr.com/reference#attribute-types"""

    # NOTE(pcattori): sphinx_autodoc_typehints cannot handle recursive reference
    # docstring written manually
    _tag: ClassVar[str] = "MAP"
    inner_type: "AttributeType"


@dataclass(frozen=True)
class Record:
    """See https://docs.tamr.com/reference#attribute-types"""

    # NOTE(pcattori) sphinx_autodoc_typehints cannot handle recursive reference
    # docstring written manually
    _tag: ClassVar[str] = "RECORD"
    attributes: Tuple[tc.SubAttribute, ...]


ComplexType = Union[Array, Map, Record]

# attribute type
################

AttributeType = Union[PrimitiveType, ComplexType]


def from_json(data: JsonDict) -> AttributeType:
    """Make an attribute type from JSON data (deserialize)

    Args:
        data: JSON data from Tamr server
    """
    base_type = data.get("baseType")
    if base_type is None:
        logger.error(f"JSON data: {repr(data)}")
        raise ValueError("Missing required field 'baseType'.")
    if base_type == Boolean._tag:
        return BOOLEAN
    elif base_type == Double._tag:
        return DOUBLE
    elif base_type == Int._tag:
        return INT
    elif base_type == Long._tag:
        return LONG
    elif base_type == String._tag:
        return STRING
    elif base_type == Array._tag:
        inner_type = data.get("innerType")
        if inner_type is None:
            logger.error(f"JSON data: {repr(data)}")
            raise ValueError("Missing required field 'innerType' for Array type.")
        return Array(inner_type=from_json(inner_type))
    elif base_type == Map._tag:
        inner_type = data.get("innerType")
        if inner_type is None:
            logger.error(f"JSON data: {repr(data)}")
            raise ValueError("Missing required field 'innerType' for Map type.")
        return Map(inner_type=from_json(inner_type))
    elif base_type == Record._tag:
        attributes = data.get("attributes")
        if attributes is None:
            logger.error(f"JSON data: {repr(data)}")
            raise ValueError("Missing required field 'attributes' for Record type.")
        return Record(
            attributes=tuple([tc.subattribute.from_json(attr) for attr in attributes])
        )
    else:
        logger.error(f"JSON data: {repr(data)}")
        raise ValueError(f"Unrecognized 'baseType': {base_type}")


def to_json(attr_type: AttributeType) -> JsonDict:
    """Serialize attribute type to JSON

    Args:
        attr_type: Attribute type to serialize
    """
    if isinstance(attr_type, (Boolean, Double, Int, Long, String)):
        return {"baseType": type(attr_type)._tag}
    elif isinstance(attr_type, (Array, Map)):
        return {
            "baseType": type(attr_type)._tag,
            "innerType": to_json(attr_type.inner_type),
        }
    elif isinstance(attr_type, Record):

        return {
            "baseType": type(attr_type)._tag,
            "attributes": [
                tc.subattribute.to_json(attr) for attr in attr_type.attributes
            ],
        }
    else:
        raise TypeError(attr_type)


# Singletons

BOOLEAN = Boolean()
DOUBLE = Double()
INT = Int()
LONG = Long()
STRING = String()