cognite-neat 0.98.0__py3-none-any.whl → 0.99.1__py3-none-any.whl

cognite/neat/_rules/models/mapping/_base.py

@@ -1,131 +0,0 @@
-from collections import Counter, defaultdict
-from collections.abc import Iterator, MutableSequence, Sequence
-from pathlib import Path
-from typing import Any, Generic, SupportsIndex, TypeVar, cast, get_args, overload
-
-import pandas as pd
-from pydantic import BaseModel, GetCoreSchemaHandler, field_validator
-from pydantic_core import core_schema
-from pydantic_core.core_schema import ValidationInfo
-
-from cognite.neat._issues.errors import NeatValueError
-from cognite.neat._rules.models._base_rules import ClassRef, PropertyRef
-from cognite.neat._rules.models.entities import ClassEntity, Undefined
-
-T_Mapping = TypeVar("T_Mapping", bound=ClassRef | PropertyRef)
-
-
-class Mapping(BaseModel, Generic[T_Mapping]):
-    source: T_Mapping
-    destination: T_Mapping
-
-
-class MappingList(list, MutableSequence[Mapping[T_Mapping]]):
-    @classmethod
-    def __get_pydantic_core_schema__(cls, source: Any, handler: GetCoreSchemaHandler) -> core_schema.CoreSchema:
-        if args := get_args(source):
-            item_type = args[0]
-        else:
-            # Someone use SheetList without specifying the type
-            raise TypeError("SheetList must be used with a type argument, e.g., SheetList[InformationProperty]")
-
-        instance_schema = core_schema.is_instance_schema(cls)
-        sequence_row_schema = handler.generate_schema(Sequence[item_type])  # type: ignore[valid-type]
-
-        non_instance_schema = core_schema.no_info_after_validator_function(MappingList, sequence_row_schema)
-        return core_schema.union_schema([instance_schema, non_instance_schema])
-
-    def as_destination_by_source(self) -> dict[T_Mapping, T_Mapping]:
-        return {mapping.source: mapping.destination for mapping in self}
-
-    def to_pandas(self, drop_na_columns: bool = True, include: list[str] | None = None) -> pd.DataFrame:
-        """Converts ResourceDict to pandas DataFrame."""
-        df = pd.DataFrame([entity.model_dump() for entity in self])
-        if drop_na_columns:
-            df = df.dropna(axis=1, how="all")
-        if include is not None:
-            df = df[include]
-        return df
-
-    def _repr_html_(self) -> str:
-        """Returns HTML representation of ResourceDict."""
-        return self.to_pandas(drop_na_columns=True)._repr_html_()  # type: ignore[operator]
-
-    # Implemented to get correct type hints
-    def __iter__(self) -> Iterator[Mapping[T_Mapping]]:
-        return super().__iter__()
-
-    @overload
-    def __getitem__(self, index: SupportsIndex) -> Mapping[T_Mapping]: ...
-
-    @overload
-    def __getitem__(self, index: slice) -> "MappingList[T_Mapping]": ...
-
-    def __getitem__(self, index: SupportsIndex | slice, /) -> "Mapping[T_Mapping] | MappingList[T_Mapping]":
-        if isinstance(index, slice):
-            return MappingList[T_Mapping](super().__getitem__(index))
-        return super().__getitem__(index)
-
-
-class RuleMapping(BaseModel):
-    properties: MappingList[PropertyRef]
-    classes: MappingList[ClassRef]
-
-    @field_validator("properties", "classes", mode="before")
-    def as_mapping_list(cls, value: Sequence[Any], info: ValidationInfo) -> Any:
-        if isinstance(value, Sequence) and not isinstance(value, MappingList):
-            annotation = cast(type, cls.model_fields[info.field_name].annotation)  # type: ignore[index]
-            ref_cls = get_args(annotation)[0]
-            return annotation([Mapping[ref_cls].model_validate(item) for item in value])  # type: ignore[valid-type, index]
-        return value
-
-    @classmethod
-    def load_spreadsheet(
-        cls, path: str | Path, source_prefix: str | None = None, destination_prefix: str | None = None
-    ) -> "RuleMapping":
-        """Loads mapping from Excel spreadsheet.
-
-        This method expects four columns in the spreadsheet. The first two columns are the source class and
-        property, and the last two columns are the destination class and property. The method will create
-        a mapping for each row in the spreadsheet.
-
-        The class mapping will be inferred from the property mappings. If a source class has multiple
-        destination classes, the most common destination class will be used.
-
-        Args:
-            path: Path to Excel spreadsheet.
-            source_prefix: Default prefix for source classes.
-            destination_prefix: Default prefix for destination classes.
-
-        Returns:
-            Mapping object.
-
-        """
-        df = pd.read_excel(path).dropna(axis=1, how="all")
-        properties = MappingList[PropertyRef]()
-        destination_classes_by_source: dict[ClassEntity, Counter[ClassEntity]] = defaultdict(Counter)
-        for _, row in df.iterrows():
-            if len(row) < 4:
-                raise NeatValueError(f"Row {row} is not valid. Expected 4 columns, got {len(row)}")
-
-            if any(pd.isna(row.iloc[:4])):
-                continue
-            source_class, source_property, destination_class, destination_property = row.iloc[:4]
-            source_entity = ClassEntity.load(source_class, prefix=source_prefix or Undefined)
-            destination_entity = ClassEntity.load(destination_class, prefix=destination_prefix or Undefined)
-            properties.append(
-                Mapping(
-                    source=PropertyRef(Class=source_entity, Property=source_property),
-                    destination=PropertyRef(Class=destination_entity, Property=destination_property),
-                )
-            )
-            destination_classes_by_source[source_entity][destination_entity] += 1
-
-        classes = MappingList[ClassRef]()
-        for source_entity, destination_classes in destination_classes_by_source.items():
-            destination_entity = destination_classes.most_common(1)[0][0]
-            classes.append(
-                Mapping(source=ClassRef(Class=source_entity), destination=ClassRef(Class=destination_entity))
-            )
-
-        return cls(properties=properties, classes=classes)

cognite/neat/_utils/cdf/loaders/__init__.py

@@ -1,25 +0,0 @@
-from ._base import ResourceLoader
-from ._data_modeling import (
-    ContainerLoader,
-    DataModelingLoader,
-    DataModelLoader,
-    SpaceLoader,
-    ViewLoader,
-)
-from ._ingestion import (
-    RawDatabaseLoader,
-    RawTableLoader,
-    TransformationLoader,
-)
-
-__all__ = [
-    "DataModelingLoader",
-    "ContainerLoader",
-    "DataModelLoader",
-    "ResourceLoader",
-    "SpaceLoader",
-    "ViewLoader",
-    "TransformationLoader",
-    "RawTableLoader",
-    "RawDatabaseLoader",
-]

cognite/neat/_utils/cdf/loaders/_base.py

@@ -1,54 +0,0 @@
-from abc import ABC, abstractmethod
-from collections.abc import Sequence
-from typing import Generic, TypeVar
-
-from cognite.client import CogniteClient
-from cognite.client.data_classes._base import (
-    T_CogniteResourceList,
-    T_WritableCogniteResource,
-    T_WriteClass,
-    WriteableCogniteResourceList,
-)
-from cognite.client.utils.useful_types import SequenceNotStr
-
-from cognite.neat._shared import T_ID
-
-T_WritableCogniteResourceList = TypeVar("T_WritableCogniteResourceList", bound=WriteableCogniteResourceList)
-
-
-class ResourceLoader(
-    ABC,
-    Generic[T_ID, T_WriteClass, T_WritableCogniteResource, T_CogniteResourceList, T_WritableCogniteResourceList],
-):
-    resource_name: str
-
-    def __init__(self, client: CogniteClient) -> None:
-        self.client = client
-
-    @classmethod
-    @abstractmethod
-    def get_id(cls, item: T_WriteClass | T_WritableCogniteResource | dict | T_ID) -> T_ID:
-        raise NotImplementedError
-
-    @classmethod
-    def get_ids(cls, items: Sequence[T_WriteClass | T_WritableCogniteResource]) -> list[T_ID]:
-        return [cls.get_id(item) for item in items]
-
-    @abstractmethod
-    def create(self, items: Sequence[T_WriteClass]) -> T_WritableCogniteResourceList:
-        raise NotImplementedError
-
-    @abstractmethod
-    def retrieve(self, ids: SequenceNotStr[T_ID]) -> T_WritableCogniteResourceList:
-        raise NotImplementedError
-
-    @abstractmethod
-    def update(self, items: Sequence[T_WriteClass]) -> T_WritableCogniteResourceList:
-        raise NotImplementedError
-
-    @abstractmethod
-    def delete(self, ids: SequenceNotStr[T_ID]) -> list[T_ID]:
-        raise NotImplementedError
-
-    def are_equal(self, local: T_WriteClass, remote: T_WritableCogniteResource) -> bool:
-        return local == remote.as_write()

cognite/neat/_utils/cdf/loaders/_data_modeling.py

@@ -1,339 +0,0 @@
-from collections.abc import Callable, Sequence
-from graphlib import TopologicalSorter
-from typing import Any, Literal, cast
-
-from cognite.client import CogniteClient
-from cognite.client.data_classes import filters
-from cognite.client.data_classes._base import (
-    T_CogniteResourceList,
-    T_WritableCogniteResource,
-    T_WriteClass,
-)
-from cognite.client.data_classes.data_modeling import (
-    Container,
-    ContainerApply,
-    ContainerApplyList,
-    ContainerList,
-    DataModel,
-    DataModelApply,
-    DataModelApplyList,
-    DataModelList,
-    RequiresConstraint,
-    Space,
-    SpaceApply,
-    SpaceApplyList,
-    SpaceList,
-    View,
-    ViewApply,
-    ViewApplyList,
-    ViewList,
-)
-from cognite.client.data_classes.data_modeling.ids import (
-    ContainerId,
-    DataModelId,
-    NodeId,
-    ViewId,
-)
-from cognite.client.exceptions import CogniteAPIError
-from cognite.client.utils.useful_types import SequenceNotStr
-
-from ._base import T_ID, ResourceLoader, T_WritableCogniteResourceList
-
-
-class DataModelingLoader(
-    ResourceLoader[T_ID, T_WriteClass, T_WritableCogniteResource, T_CogniteResourceList, T_WritableCogniteResourceList]
-):
-    @classmethod
-    def in_space(cls, item: T_WriteClass | T_WritableCogniteResource | T_ID, space: set[str]) -> bool:
-        if hasattr(item, "space"):
-            return item.space in space
-        raise ValueError(f"Item {item} does not have a space attribute")
-
-    def sort_by_dependencies(self, items: list[T_WriteClass]) -> list[T_WriteClass]:
-        return items
-
-    def _create_force(
-        self,
-        items: Sequence[T_WriteClass],
-        tried_force_deploy: set[T_ID],
-        create_method: Callable[[Sequence[T_WriteClass]], T_WritableCogniteResourceList],
-    ) -> T_WritableCogniteResourceList:
-        try:
-            return create_method(items)
-        except CogniteAPIError as e:
-            failed_items = {failed.as_id() for failed in e.failed if hasattr(failed, "as_id")}
-            to_redeploy = [
-                item
-                for item in items
-                if item.as_id() in failed_items and item.as_id() not in tried_force_deploy  # type: ignore[attr-defined]
-            ]
-            if not to_redeploy:
-                # Avoid infinite loop
-                raise e
-            ids = [item.as_id() for item in to_redeploy]  # type: ignore[attr-defined]
-            tried_force_deploy.update(ids)
-            self.delete(ids)
-            return self._create_force(to_redeploy, tried_force_deploy, create_method)
-
-
-class SpaceLoader(DataModelingLoader[str, SpaceApply, Space, SpaceApplyList, SpaceList]):
-    resource_name = "spaces"
-
-    @classmethod
-    def get_id(cls, item: Space | SpaceApply | str | dict) -> str:
-        if isinstance(item, Space | SpaceApply):
-            return item.space
-        if isinstance(item, dict):
-            return item["space"]
-        return item
-
-    def create(self, items: Sequence[SpaceApply]) -> SpaceList:
-        return self.client.data_modeling.spaces.apply(items)
-
-    def retrieve(self, ids: SequenceNotStr[str]) -> SpaceList:
-        return self.client.data_modeling.spaces.retrieve(ids)
-
-    def update(self, items: Sequence[SpaceApply]) -> SpaceList:
-        return self.create(items)
-
-    def delete(self, ids: SequenceNotStr[str] | Sequence[Space | SpaceApply]) -> list[str]:
-        if all(isinstance(item, Space) for item in ids) or all(isinstance(item, SpaceApply) for item in ids):
-            ids = [cast(Space | SpaceApply, item).space for item in ids]
-        return self.client.data_modeling.spaces.delete(cast(SequenceNotStr[str], ids))
-
-    def clean(self, space: str) -> None:
-        """Deletes all data in a space.
-
-        This means all nodes, edges, views, containers, and data models located in the given space.
-
-        Args:
-            client: Connected CogniteClient
-            space: The space to delete.
-
-        """
-        edges = self.client.data_modeling.instances.list(
-            "edge", limit=-1, filter=filters.Equals(["edge", "space"], space)
-        )
-        if edges:
-            instances = self.client.data_modeling.instances.delete(edges=edges.as_ids())
-            print(f"Deleted {len(instances.edges)} edges")
-        nodes = self.client.data_modeling.instances.list(
-            "node", limit=-1, filter=filters.Equals(["node", "space"], space)
-        )
-        node_types = {NodeId(node.type.space, node.type.external_id) for node in nodes if node.type}
-        node_data = set(nodes.as_ids()) - node_types
-        if node_data:
-            instances = self.client.data_modeling.instances.delete(nodes=list(node_data))
-            print(f"Deleted {len(instances.nodes)} nodes")
-        if node_types:
-            instances = self.client.data_modeling.instances.delete(nodes=list(node_types))
-            print(f"Deleted {len(instances.nodes)} node types")
-        views = self.client.data_modeling.views.list(limit=-1, space=space)
-        if views:
-            deleted_views = self.client.data_modeling.views.delete(views.as_ids())
-            print(f"Deleted {len(deleted_views)} views")
-        containers = self.client.data_modeling.containers.list(limit=-1, space=space)
-        if containers:
-            deleted_containers = self.client.data_modeling.containers.delete(containers.as_ids())
-            print(f"Deleted {len(deleted_containers)} containers")
-        if data_models := self.client.data_modeling.data_models.list(limit=-1, space=space):
-            deleted_data_models = self.client.data_modeling.data_models.delete(data_models.as_ids())
-            print(f"Deleted {len(deleted_data_models)} data models")
-        deleted_space = self.client.data_modeling.spaces.delete(space)
-        print(f"Deleted space {deleted_space}")
-
-
-class ViewLoader(DataModelingLoader[ViewId, ViewApply, View, ViewApplyList, ViewList]):
-    resource_name = "views"
-
-    def __init__(self, client: CogniteClient, existing_handling: Literal["fail", "skip", "update", "force"] = "fail"):
-        super().__init__(client)
-        self.existing_handling = existing_handling
-        self._cache_view_by_id: dict[ViewId, View] = {}
-        self._tried_force_deploy: set[ViewId] = set()
-
-    @classmethod
-    def get_id(cls, item: View | ViewApply | ViewId | dict) -> ViewId:
-        if isinstance(item, View | ViewApply):
-            return item.as_id()
-        if isinstance(item, dict):
-            return ViewId.load(item)
-        return item
-
-    def create(self, items: Sequence[ViewApply]) -> ViewList:
-        if self.existing_handling == "force":
-            return self._create_force(items, self._tried_force_deploy, self.client.data_modeling.views.apply)
-        else:
-            return self.client.data_modeling.views.apply(items)
-
-    def retrieve(self, ids: SequenceNotStr[ViewId]) -> ViewList:
-        return self.client.data_modeling.views.retrieve(cast(Sequence, ids))
-
-    def update(self, items: Sequence[ViewApply]) -> ViewList:
-        return self.create(items)
-
-    def delete(self, ids: SequenceNotStr[ViewId]) -> list[ViewId]:
-        return self.client.data_modeling.views.delete(cast(Sequence, ids))
-
-    def _as_write_raw(self, view: View) -> dict[str, Any]:
-        dumped = view.as_write().dump()
-        if view.properties:
-            # All read version of views have all the properties of their parent views.
-            # We need to remove these properties to compare with the local view.
-            parents = self._retrieve_view_ancestors(view.implements or [], self._cache_view_by_id)
-            for parent in parents:
-                for prop_name in parent.properties.keys():
-                    dumped["properties"].pop(prop_name, None)
-
-        if "properties" in dumped and not dumped["properties"]:
-            # All properties were removed, so we remove the properties key.
-            dumped.pop("properties", None)
-        return dumped
-
-    def are_equal(self, local: ViewApply, remote: View) -> bool:
-        local_dumped = local.dump()
-        if not remote.implements:
-            return local_dumped == remote.as_write().dump()
-
-        cdf_resource_dumped = self._as_write_raw(remote)
-
-        if "properties" in local_dumped and not local_dumped["properties"]:
-            # In case the local properties are set to an empty dict.
-            local_dumped.pop("properties", None)
-
-        return local_dumped == cdf_resource_dumped
-
-    def as_write(self, view: View) -> ViewApply:
-        return ViewApply.load(self._as_write_raw(view))
-
-    def retrieve_all_parents(self, views: list[ViewId]) -> list[View]:
-        return self._retrieve_view_ancestors(views, self._cache_view_by_id)
-
-    def _retrieve_view_ancestors(self, parents: list[ViewId], cache: dict[ViewId, View]) -> list[View]:
-        """Retrieves all ancestors of a view.
-
-        This will mutate the cache passed in, and return a list of views that are the ancestors
-        of the views in the parents list.
-
-        Args:
-            parents: The parents of the view to retrieve all ancestors for
-            cache: The cache to store the views in
-        """
-        parent_ids = parents.copy()
-        found: list[View] = []
-        found_ids: set[ViewId] = set()
-        while parent_ids:
-            to_lookup: set[ViewId] = set()
-            grand_parent_ids = []
-            for parent in parent_ids:
-                if parent in found_ids:
-                    continue
-                elif parent in cache:
-                    found.append(cache[parent])
-                    grand_parent_ids.extend(cache[parent].implements or [])
-                else:
-                    to_lookup.add(parent)
-
-            if to_lookup:
-                looked_up = self.client.data_modeling.views.retrieve(list(to_lookup))
-                cache.update({view.as_id(): view for view in looked_up})
-                found.extend(looked_up)
-                found_ids.update({view.as_id() for view in looked_up})
-                for view in looked_up:
-                    grand_parent_ids.extend(view.implements or [])
-
-            parent_ids = grand_parent_ids
-        return found
-
-
-class ContainerLoader(DataModelingLoader[ContainerId, ContainerApply, Container, ContainerApplyList, ContainerList]):
-    resource_name = "containers"
-
-    def __init__(self, client: CogniteClient, existing_handling: Literal["fail", "skip", "update", "force"] = "fail"):
-        super().__init__(client)
-        self.existing_handling = existing_handling
-        self._tried_force_deploy: set[ContainerId] = set()
-
-    @classmethod
-    def get_id(cls, item: Container | ContainerApply | ContainerId | dict) -> ContainerId:
-        if isinstance(item, Container | ContainerApply):
-            return item.as_id()
-        if isinstance(item, dict):
-            return ContainerId.load(item)
-        return item
-
-    def sort_by_dependencies(self, items: Sequence[ContainerApply]) -> list[ContainerApply]:
-        container_by_id = {container.as_id(): container for container in items}
-        container_dependencies = {
-            container.as_id(): {
-                const.require
-                for const in container.constraints.values()
-                if isinstance(const, RequiresConstraint) and const.require in container_by_id
-            }
-            for container in items
-        }
-        return [
-            container_by_id[container_id] for container_id in TopologicalSorter(container_dependencies).static_order()
-        ]
-
-    def create(self, items: Sequence[ContainerApply]) -> ContainerList:
-        if self.existing_handling == "force":
-            return self._create_force(items, self._tried_force_deploy, self.client.data_modeling.containers.apply)
-        else:
-            return self.client.data_modeling.containers.apply(items)
-
-    def retrieve(self, ids: SequenceNotStr[ContainerId]) -> ContainerList:
-        return self.client.data_modeling.containers.retrieve(cast(Sequence, ids))
-
-    def update(self, items: Sequence[ContainerApply]) -> ContainerList:
-        return self.create(items)
-
-    def delete(self, ids: SequenceNotStr[ContainerId]) -> list[ContainerId]:
-        return self.client.data_modeling.containers.delete(cast(Sequence, ids))
-
-    def are_equal(self, local: ContainerApply, remote: Container) -> bool:
-        local_dumped = local.dump(camel_case=True)
-        if "usedFor" not in local_dumped:
-            # Setting used_for to "node" as it is the default value in the CDF.
-            local_dumped["usedFor"] = "node"
-
-        return local_dumped == remote.as_write().dump(camel_case=True)
-
-
-class DataModelLoader(DataModelingLoader[DataModelId, DataModelApply, DataModel, DataModelApplyList, DataModelList]):
-    resource_name = "data_models"
-
-    @classmethod
-    def get_id(cls, item: DataModel | DataModelApply | DataModelId | dict) -> DataModelId:
-        if isinstance(item, DataModel | DataModelApply):
-            return item.as_id()
-        if isinstance(item, dict):
-            return DataModelId.load(item)
-        return item
-
-    def create(self, items: Sequence[DataModelApply]) -> DataModelList:
-        return self.client.data_modeling.data_models.apply(items)
-
-    def retrieve(self, ids: SequenceNotStr[DataModelId]) -> DataModelList:
-        return self.client.data_modeling.data_models.retrieve(cast(Sequence, ids))
-
-    def update(self, items: Sequence[DataModelApply]) -> DataModelList:
-        return self.create(items)
-
-    def delete(self, ids: SequenceNotStr[DataModelId]) -> list[DataModelId]:
-        return self.client.data_modeling.data_models.delete(cast(Sequence, ids))
-
-    def are_equal(self, local: DataModelApply, remote: DataModel) -> bool:
-        local_dumped = local.dump()
-        cdf_resource_dumped = remote.as_write().dump()
-
-        # Data models that have the same views, but in different order, are considered equal.
-        # We also account for whether views are given as IDs or View objects.
-        local_dumped["views"] = sorted(
-            (v if isinstance(v, ViewId) else v.as_id()).as_tuple() for v in local.views or []
-        )
-        cdf_resource_dumped["views"] = sorted(
-            (v if isinstance(v, ViewId) else v.as_id()).as_tuple() for v in remote.views or []
-        )
-
-        return local_dumped == cdf_resource_dumped