pytrilogy 0.3.142__cp312-cp312-win_amd64.whl

trilogy/core/models/build_environment.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+import difflib
+from pathlib import Path
+from typing import Annotated, Dict, ItemsView, Never, ValuesView
+
+from pydantic import BaseModel, ConfigDict, Field
+from pydantic.functional_validators import PlainValidator
+
+from trilogy.constants import DEFAULT_NAMESPACE
+from trilogy.core.exceptions import (
+    UndefinedConceptException,
+)
+from trilogy.core.models.build import BuildConcept, BuildDatasource, BuildFunction
+from trilogy.core.models.core import DataType
+
+
+class BuildEnvironmentConceptDict(dict):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(self, *args, **kwargs)
+
+    def values(self) -> ValuesView[BuildConcept]:  # type: ignore
+        return super().values()
+
+    def get(self, key: str, default: BuildConcept | None = None) -> BuildConcept | None:  # type: ignore
+        try:
+            return self.__getitem__(key)
+        except UndefinedConceptException:
+            return default
+
+    def raise_undefined(
+        self, key: str, line_no: int | None = None, file: Path | str | None = None
+    ) -> Never:
+        # build environment should never check for missing values.
+        if line_no is not None:
+            message = f"Concept '{key}' not found in environment at line {line_no}."
+        else:
+            message = f"Concept '{key}' not found in environment."
+        raise UndefinedConceptException(message, [])
+
+    def __getitem__(
+        self, key: str, line_no: int | None = None, file: Path | None = None
+    ) -> BuildConcept:
+        try:
+            return super(BuildEnvironmentConceptDict, self).__getitem__(key)
+        except KeyError:
+            if "." in key and key.split(".", 1)[0] == DEFAULT_NAMESPACE:
+                return self.__getitem__(key.split(".", 1)[1], line_no)
+            if DEFAULT_NAMESPACE + "." + key in self:
+                return self.__getitem__(DEFAULT_NAMESPACE + "." + key, line_no)
+        self.raise_undefined(key, line_no, file)
+
+    def _find_similar_concepts(self, concept_name: str):
+        def strip_local(input: str):
+            if input.startswith(f"{DEFAULT_NAMESPACE}."):
+                return input[len(DEFAULT_NAMESPACE) + 1 :]
+            return input
+
+        matches = difflib.get_close_matches(
+            strip_local(concept_name), [strip_local(x) for x in self.keys()]
+        )
+        return matches
+
+    def items(self) -> ItemsView[str, BuildConcept]:  # type: ignore
+        return super().items()
+
+
+class BuildEnvironmentDatasourceDict(dict):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(self, *args, **kwargs)
+
+    def __getitem__(self, key: str) -> BuildDatasource:
+        try:
+            return super(BuildEnvironmentDatasourceDict, self).__getitem__(key)
+        except KeyError:
+            if DEFAULT_NAMESPACE + "." + key in self:
+                return self.__getitem__(DEFAULT_NAMESPACE + "." + key)
+            if "." in key and key.split(".", 1)[0] == DEFAULT_NAMESPACE:
+                return self.__getitem__(key.split(".", 1)[1])
+            raise
+
+    def values(self) -> ValuesView[BuildDatasource]:  # type: ignore
+        return super().values()
+
+    def items(self) -> ItemsView[str, BuildDatasource]:  # type: ignore
+        return super().items()
+
+
+def validate_concepts(v) -> BuildEnvironmentConceptDict:
+    if isinstance(v, BuildEnvironmentConceptDict):
+        return v
+    elif isinstance(v, dict):
+        return BuildEnvironmentConceptDict(**{x: y for x, y in v.items()})
+    raise ValueError
+
+
+def validate_datasources(v) -> BuildEnvironmentDatasourceDict:
+    if isinstance(v, BuildEnvironmentDatasourceDict):
+        return v
+    elif isinstance(v, dict):
+        return BuildEnvironmentDatasourceDict(**{x: y for x, y in v.items()})
+    raise ValueError
+
+
+class BuildEnvironment(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True, strict=False)
+
+    concepts: Annotated[
+        BuildEnvironmentConceptDict, PlainValidator(validate_concepts)
+    ] = Field(default_factory=BuildEnvironmentConceptDict)
+
+    canonical_concepts: Annotated[
+        BuildEnvironmentConceptDict, PlainValidator(validate_concepts)
+    ] = Field(default_factory=BuildEnvironmentConceptDict)
+
+    datasources: Annotated[
+        BuildEnvironmentDatasourceDict, PlainValidator(validate_datasources)
+    ] = Field(default_factory=BuildEnvironmentDatasourceDict)
+    functions: Dict[str, BuildFunction] = Field(default_factory=dict)
+    data_types: Dict[str, DataType] = Field(default_factory=dict)
+    namespace: str = DEFAULT_NAMESPACE
+    cte_name_map: Dict[str, str] = Field(default_factory=dict)
+    materialized_concepts: set[str] = Field(default_factory=set)
+    materialized_canonical_concepts: set[str] = Field(default_factory=set)
+    non_partial_materialized_canonical_concepts: set[str] = Field(default_factory=set)
+    alias_origin_lookup: Dict[str, BuildConcept] = Field(default_factory=dict)
+
+    def gen_concept_list_caches(self) -> None:
+        concrete_concepts: list[BuildConcept] = []
+        non_partial_concrete_concepts: list[BuildConcept] = []
+        for datasource in self.datasources.values():
+            for column in datasource.columns:
+                if column.is_complete:
+                    non_partial_concrete_concepts.append(column.concept)
+                concrete_concepts.append(column.concept)
+        concrete_addresses = set([x.address for x in concrete_concepts])
+        canonical_addresses = set([x.canonical_address for x in concrete_concepts])
+        non_partial_canonical_addresses = set(
+            [x.canonical_address for x in non_partial_concrete_concepts]
+        )
+
+        self.materialized_concepts = set(
+            [
+                c.address
+                for c in self.concepts.values()
+                if c.address in concrete_addresses
+            ]
+            + [
+                c.address
+                for c in self.alias_origin_lookup.values()
+                if c.address in concrete_addresses
+            ],
+        )
+        self.materialized_canonical_concepts = set(
+            [
+                c.canonical_address
+                for c in self.concepts.values()
+                if c.canonical_address in canonical_addresses
+            ]
+            + [
+                c.canonical_address
+                for c in self.alias_origin_lookup.values()
+                if c.canonical_address in canonical_addresses
+            ],
+        )
+        self.non_partial_materialized_canonical_concepts = set(
+            [
+                c.canonical_address
+                for c in self.concepts.values()
+                if c.canonical_address in non_partial_canonical_addresses
+            ]
+            + [
+                c.canonical_address
+                for c in self.alias_origin_lookup.values()
+                if c.canonical_address in non_partial_canonical_addresses
+            ],
+        )
+
+
+BuildEnvironment.model_rebuild()

trilogy/core/models/core.py

@@ -0,0 +1,501 @@
+from __future__ import annotations
+
+from abc import ABC
+from collections import UserDict, UserList
+from datetime import date, datetime
+from decimal import Decimal
+from enum import Enum
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Generic,
+    List,
+    Sequence,
+    Tuple,
+    TypeVar,
+    Union,
+    get_args,
+)
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    field_validator,
+)
+from pydantic_core import core_schema
+
+from trilogy.constants import (
+    MagicConstants,
+)
+from trilogy.core.enums import DatePart, Modifier, Ordering
+
+
+class DataTyped(ABC):
+
+    # this is not abstract
+    # only because when it's a pydantic property, it fails validation
+    @property
+    def output_datatype(self) -> CONCRETE_TYPES:  # type: ignore
+        """
+        This is a huge hack to get property vs pydantic attribute inheritance to work.
+        """
+        if "output_datatype" in self.__dict__:
+            return self.__dict__["output_datatype"]
+        raise NotImplementedError
+
+
+class Addressable(ABC):
+
+    @property
+    def _address(self):
+        return self.address
+
+
+TYPEDEF_TYPES = Union[
+    "DataType",
+    "MapType",
+    "ArrayType",
+    "NumericType",
+    "StructType",
+    "DataTyped",
+    "TraitDataType",
+]
+
+CONCRETE_TYPES = Union[
+    "DataType",
+    "MapType",
+    "ArrayType",
+    "NumericType",
+    "StructType",
+    "TraitDataType",
+]
+
+KT = TypeVar("KT")
+VT = TypeVar("VT")
+LT = TypeVar("LT")
+
+
+class DataType(Enum):
+    # PRIMITIVES
+    STRING = "string"
+    BOOL = "bool"
+    MAP = "map"
+    NUMBER = "number"
+    FLOAT = "float"
+    NUMERIC = "numeric"
+    INTEGER = "int"
+    BIGINT = "bigint"
+    DATE = "date"
+    DATETIME = "datetime"
+    TIMESTAMP = "timestamp"
+    ARRAY = "array"
+    DATE_PART = "date_part"
+    STRUCT = "struct"
+    GEOGRAPHY = "geography"
+    NULL = "null"
+
+    # GRANULAR
+    UNIX_SECONDS = "unix_seconds"
+
+    # PARSING
+    UNKNOWN = "unknown"
+    ANY = "any"
+
+    @property
+    def data_type(self):
+        return self
+
+    def __str__(self) -> str:
+        return self.name
+
+
+class TraitDataType(BaseModel):
+    type: DataType | NumericType | StructType | ArrayType | MapType
+    traits: list[str]
+
+    def __hash__(self):
+        return hash(self.type)
+
+    def __str__(self) -> str:
+        return f"Trait<{self.type}, {self.traits}>"
+
+    def __eq__(self, other):
+        if isinstance(other, DataType):
+            return self.type == other
+        elif isinstance(other, TraitDataType):
+            return self.type == other.type and self.traits == other.traits
+        return False
+
+    @property
+    def data_type(self):
+        return self.type
+
+    @property
+    def value(self):
+        return self.data_type.value
+
+
+class NumericType(BaseModel):
+    precision: int = 20
+    scale: int = 5
+
+    def __str__(self) -> str:
+        return f"Numeric({self.precision},{self.scale})"
+
+    def __hash__(self):
+        return hash((DataType.NUMERIC, self.precision, self.scale))
+
+    @property
+    def data_type(self):
+        return DataType.NUMERIC
+
+    @property
+    def value(self):
+        return self.data_type.value
+
+
+class ArrayType(BaseModel):
+    model_config = ConfigDict(frozen=True)
+    type: TYPEDEF_TYPES
+
+    @field_validator("type", mode="plain")
+    def validate_type(cls, v):
+        return v
+
+    def __str__(self) -> str:
+        return f"ListType<{self.type}>"
+
+    @property
+    def data_type(self):
+        return DataType.ARRAY
+
+    @property
+    def value(self) -> str:
+        return self.data_type.value
+
+    @property
+    def value_data_type(
+        self,
+    ) -> CONCRETE_TYPES:
+        if isinstance(self.type, DataTyped):
+            return self.type.output_datatype
+        return self.type
+
+
+class MapType(BaseModel):
+    key_type: TYPEDEF_TYPES
+    value_type: TYPEDEF_TYPES
+
+    @field_validator("value_type", mode="plain")
+    def validate_type(cls, v):
+        return v
+
+    @field_validator("key_type", mode="plain")
+    def validate_key_type(cls, v):
+        return v
+
+    def __hash__(self):
+        return hash((DataType.MAP, self.key_type, self.value_type))
+
+    @property
+    def data_type(self):
+        return DataType.MAP
+
+    @property
+    def value(self):
+        return self.data_type.value
+
+    @property
+    def value_data_type(
+        self,
+    ) -> CONCRETE_TYPES:
+        if isinstance(self.value_type, DataTyped):
+            return self.value_type.output_datatype
+        return self.value_type
+
+    @property
+    def key_data_type(
+        self,
+    ) -> CONCRETE_TYPES:
+        if isinstance(self.key_type, DataTyped):
+            return self.key_type.output_datatype
+        return self.key_type
+
+
+class StructComponent(BaseModel):
+    name: str
+    type: TYPEDEF_TYPES
+    modifiers: list[Modifier] = []
+
+    @field_validator("type", mode="plain")
+    def validate_type(cls, v):
+        return v
+
+
+class StructType(BaseModel):
+    fields: Sequence[StructComponent | TYPEDEF_TYPES]
+    fields_map: Dict[str, DataTyped | int | float | str | StructComponent]
+
+    def __repr__(self):
+        return "struct<{}>".format(
+            ", ".join(
+                f"{field.name}:{field.type.name}"
+                for field in self.fields
+                if isinstance(field, StructComponent)
+            )
+        )
+
+    def __str__(self) -> str:
+        return self.__repr__()
+
+    @field_validator("fields", mode="plain")
+    def validate_type(cls, v):
+        final = []
+        for field in v:
+            final.append(field)
+        return final
+
+    @field_validator("fields_map", mode="plain")
+    def validate_fields_map(cls, v):
+        return v
+
+    @property
+    def data_type(self):
+        return DataType.STRUCT
+
+    @property
+    def value(self):
+        return self.data_type.value
+
+    @property
+    def field_types(self) -> Dict[str, CONCRETE_TYPES]:
+        out: Dict[str, CONCRETE_TYPES] = {}
+        keys = list(self.fields_map.keys())
+        for idx, field in enumerate(self.fields):
+            if isinstance(field, StructComponent):
+                out[field.name] = arg_to_datatype(field.type)
+            elif isinstance(field, DataTyped):
+                out[keys[idx]] = field.output_datatype
+            else:
+                out[keys[idx]] = field
+        return out
+
+    def __hash__(self):
+        return hash(str(self))
+
+
+class ListWrapper(Generic[VT], UserList):
+    """Used to distinguish parsed list objects from other lists"""
+
+    def __init__(self, *args, type: DataType, nullable: bool = False, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.type = type
+        self.nullable = nullable
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: Callable[[Any], core_schema.CoreSchema]
+    ) -> core_schema.CoreSchema:
+        args = get_args(source_type)
+        if args:
+            schema = handler(List[args])  # type: ignore
+        else:
+            schema = handler(List)
+        return core_schema.no_info_after_validator_function(cls.validate, schema)
+
+    @classmethod
+    def validate(cls, v):
+        return cls(v, type=arg_to_datatype(v[0]))
+
+
+class MapWrapper(Generic[KT, VT], UserDict):
+    """Used to distinguish parsed map objects from other dicts"""
+
+    def __init__(self, *args, key_type: DataType, value_type: DataType, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.key_type = key_type
+        self.value_type = value_type
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: Callable[[Any], core_schema.CoreSchema]
+    ) -> core_schema.CoreSchema:
+        args = get_args(source_type)
+        if args:
+            schema = handler(Dict[args])  # type: ignore
+        else:
+            schema = handler(Dict)
+        return core_schema.no_info_after_validator_function(cls.validate, schema)
+
+    @classmethod
+    def validate(cls, v):
+        return cls(
+            v,
+            key_type=arg_to_datatype(list(v.keys()).pop()),
+            value_type=arg_to_datatype(list(v.values()).pop()),
+        )
+
+
+class TupleWrapper(Generic[VT], tuple):
+    """Used to distinguish parsed tuple objects from other tuples"""
+
+    def __init__(self, val, type: DataType, nullable: bool = False, **kwargs):
+        super().__init__()
+        self.type = type
+        self.val = val
+        self.nullable = nullable
+
+    def __getnewargs__(self):
+        return (self.val, self.type)
+
+    def __new__(cls, val, type: DataType, **kwargs):
+        return super().__new__(cls, tuple(val))
+        # self.type = type
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: Callable[[Any], core_schema.CoreSchema]
+    ) -> core_schema.CoreSchema:
+        args = get_args(source_type)
+        if args:
+            schema = handler(Tuple[args])  # type: ignore
+        else:
+            schema = handler(Tuple)
+        return core_schema.no_info_after_validator_function(cls.validate, schema)
+
+    @classmethod
+    def validate(cls, v):
+        return cls(v, type=arg_to_datatype(v[0]))
+
+
+def list_to_wrapper(args):
+    rtypes = [arg_to_datatype(arg) for arg in args]
+    types = [arg for arg in rtypes if arg != DataType.NULL]
+    if not len(set(types)) == 1:
+        raise SyntaxError(f"Cannot create a list with this set of types: {set(types)}")
+    return ListWrapper(args, type=types[0], nullable=DataType.NULL in rtypes)
+
+
+def tuple_to_wrapper(args):
+    rtypes = [arg_to_datatype(arg) for arg in args]
+    types = [arg for arg in rtypes if arg != DataType.NULL]
+    if not len(set(types)) == 1:
+        raise SyntaxError(f"Cannot create a tuple with this set of types: {set(types)}")
+    return TupleWrapper(args, type=types[0], nullable=DataType.NULL in rtypes)
+
+
+def dict_to_map_wrapper(arg):
+    key_types = [arg_to_datatype(arg) for arg in arg.keys()]
+
+    value_types = [arg_to_datatype(arg) for arg in arg.values()]
+    assert len(set(key_types)) == 1
+    assert len(set(key_types)) == 1
+    return MapWrapper(arg, key_type=key_types[0], value_type=value_types[0])
+
+
+def merge_datatypes(
+    inputs: list[
+        DataType | ArrayType | StructType | MapType | NumericType | TraitDataType
+    ],
+) -> DataType | ArrayType | StructType | MapType | NumericType | TraitDataType:
+    """This is a temporary hack for doing between
+    allowable datatype transformation matrix"""
+    if len(inputs) == 1:
+        return inputs[0]
+
+    if set(inputs) == {DataType.INTEGER, DataType.FLOAT}:
+        return DataType.FLOAT
+    if set(inputs) == {DataType.INTEGER, DataType.NUMERIC}:
+        return DataType.NUMERIC
+    if any(isinstance(x, NumericType) for x in inputs) and all(
+        isinstance(x, NumericType)
+        or x in (DataType.INTEGER, DataType.FLOAT, DataType.NUMERIC)
+        for x in inputs
+    ):
+        candidate = next(x for x in inputs if isinstance(x, NumericType))
+        return candidate
+    return inputs[0]
+
+
+def is_compatible_datatype(left, right):
+    # for unknown types, we can't make any assumptions
+    if isinstance(left, list):
+        return any(is_compatible_datatype(ltype, right) for ltype in left)
+    if isinstance(right, list):
+        return any(is_compatible_datatype(left, rtype) for rtype in right)
+    if left == DataType.ANY or right == DataType.ANY:
+        return True
+    if all(
+        isinstance(x, NumericType)
+        or x in (DataType.INTEGER, DataType.FLOAT, DataType.NUMERIC)
+        for x in (left, right)
+    ):
+        return True
+    elif isinstance(left, NumericType) or isinstance(right, NumericType):
+        return False
+    if right == DataType.UNKNOWN or left == DataType.UNKNOWN:
+        return True
+    if left == right:
+        return True
+    if {left, right} == {DataType.NUMERIC, DataType.FLOAT}:
+        return True
+
+    if {left, right} == {DataType.NUMERIC, DataType.INTEGER}:
+        return True
+    if {left, right} == {DataType.FLOAT, DataType.INTEGER}:
+        return True
+    return False
+
+
+def arg_to_datatype(arg) -> CONCRETE_TYPES:
+
+    if isinstance(arg, MagicConstants):
+        if arg == MagicConstants.NULL:
+            return DataType.NULL
+        raise ValueError(f"Cannot parse arg datatype for arg of type {arg}")
+    elif isinstance(arg, bool):
+        return DataType.BOOL
+    elif isinstance(arg, Ordering):
+        return DataType.STRING  # TODO: revisit
+    elif isinstance(arg, int):
+        return DataType.INTEGER
+    elif isinstance(arg, str):
+        return DataType.STRING
+    elif isinstance(arg, float):
+        return DataType.FLOAT
+    elif isinstance(arg, Decimal):
+        return DataType.NUMERIC
+    elif isinstance(arg, DataType):
+        return arg
+    elif isinstance(arg, NumericType):
+        return arg
+    elif isinstance(arg, TraitDataType):
+        return arg
+    elif isinstance(arg, ListWrapper):
+        return ArrayType(type=arg.type)
+    elif isinstance(arg, ArrayType):
+        return arg
+    elif isinstance(arg, MapType):
+        return arg
+    elif isinstance(arg, DataTyped):
+        return arg.output_datatype
+    elif isinstance(arg, TupleWrapper):
+        return ArrayType(type=arg.type)
+    elif isinstance(arg, list):
+        wrapper = list_to_wrapper(arg)
+        return ArrayType(type=wrapper.type)
+    elif isinstance(arg, MapWrapper):
+        return MapType(key_type=arg.key_type, value_type=arg.value_type)
+    elif isinstance(arg, datetime):
+        return DataType.DATETIME
+    elif isinstance(arg, date):
+        return DataType.DATE
+    elif isinstance(arg, StructComponent):
+        return arg_to_datatype(arg.type)
+    elif isinstance(arg, DatePart):
+        return DataType.DATE_PART
+    else:
+        raise ValueError(
+            f"Cannot parse arg datatype for arg of raw type {type(arg)} value {arg}"
+        )