cognite-neat 0.125.1__tar.gz → 0.126.1__tar.gz

{cognite_neat-0.125.1 → cognite_neat-0.126.1}/.github/workflows/code-quality.yaml

@@ -69,7 +69,7 @@ jobs:
       - run: uv sync --all-extras
 
       - name: Run unit tests
-        run: export PYTHONPATH=$PYTHONPATH:$(pwd); uv run pytest tests/tests_unit tests/tests_unit_v1
+        run: export PYTHONPATH=$PYTHONPATH:$(pwd); uv run pytest tests/tests_unit
 
   build-docs:
     runs-on: ubuntu-latest

{cognite_neat-0.125.1 → cognite_neat-0.126.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognite-neat
-Version: 0.125.1
+Version: 0.126.1
 Summary: Knowledge graph transformation
 Project-URL: Documentation, https://cognite-neat.readthedocs-hosted.com/
 Project-URL: Homepage, https://cognite-neat.readthedocs-hosted.com/

cognite_neat-0.126.1/cognite/neat/_client/__init__.py

@@ -0,0 +1,4 @@
+from .client import NeatClient
+from .config import NeatClientConfig
+
+__all__ = ["NeatClient", "NeatClientConfig"]

cognite_neat-0.126.1/cognite/neat/_client/api.py

@@ -0,0 +1,8 @@
+from cognite.neat._client.config import NeatClientConfig
+from cognite.neat._utils.http_client import HTTPClient
+
+
+class NeatAPI:
+    def __init__(self, neat_config: NeatClientConfig, http_client: HTTPClient) -> None:
+        self._config = neat_config
+        self._http_client = http_client

cognite_neat-0.126.1/cognite/neat/_client/client.py

@@ -0,0 +1,19 @@
+from cognite.client import ClientConfig, CogniteClient
+
+from cognite.neat._utils.http_client import HTTPClient
+
+from .config import NeatClientConfig
+from .containers_api import ContainersAPI
+from .data_model_api import DataModelsAPI
+from .spaces_api import SpacesAPI
+from .views_api import ViewsAPI
+
+
+class NeatClient:
+    def __init__(self, cognite_client_or_config: CogniteClient | ClientConfig) -> None:
+        self.config = NeatClientConfig(cognite_client_or_config)
+        http_client = HTTPClient(self.config)
+        self.data_models = DataModelsAPI(self.config, http_client)
+        self.views = ViewsAPI(self.config, http_client)
+        self.containers = ContainersAPI(self.config, http_client)
+        self.spaces = SpacesAPI(self.config, http_client)

cognite_neat-0.126.1/cognite/neat/_client/config.py

@@ -0,0 +1,40 @@
+from cognite.client import ClientConfig, CogniteClient
+
+
+class NeatClientConfig(ClientConfig):
+    def __init__(self, cognite_client_or_config: CogniteClient | ClientConfig):
+        config = (
+            cognite_client_or_config.config
+            if isinstance(cognite_client_or_config, CogniteClient)
+            else cognite_client_or_config
+        )
+        super().__init__(
+            client_name=config.client_name,
+            project=config.project,
+            credentials=config.credentials,
+            api_subversion=config.api_subversion,
+            base_url=config.base_url,
+            max_workers=config.max_workers,
+            headers=config.headers,
+            timeout=config.timeout,
+            file_transfer_timeout=config.file_transfer_timeout,
+            debug=config.debug,
+        )
+
+    def create_api_url(self, endpoint: str) -> str:
+        """Create a full API URL for the given endpoint.
+
+        Args:
+            endpoint (str): The API endpoint to append to the base URL.
+
+        Returns:
+            str: The full API URL.
+
+        Examples:
+            >>> config = NeatClientConfig(cluster="bluefield", project="my_project", ...)
+            >>> config.create_api_url("/models/instances")
+            "https://bluefield.cognitedata.com/api/v1/my_project/models/instances"
+        """
+        if not endpoint.startswith("/"):
+            endpoint = f"/{endpoint}"
+        return f"{self.base_url}/api/v1/projects/{self.project}{endpoint}"

cognite_neat-0.126.1/cognite/neat/_client/containers_api.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from cognite.neat._data_model.models.dms import ContainerReference, ContainerResponse
+from cognite.neat._utils.http_client import ItemIDBody, ItemsRequest, ParametersRequest
+from cognite.neat._utils.useful_types import PrimitiveType
+
+from .api import NeatAPI
+from .data_classes import PagedResponse
+
+
+class ContainersAPI(NeatAPI):
+    def retrieve(
+        self,
+        items: list[ContainerReference],
+    ) -> list[ContainerResponse]:
+        """Retrieve containers by their identifiers.
+
+        Args:
+            items: List of (space, external_id) tuples identifying the containers to retrieve.
+
+        Returns:
+            List of ContainerResponse objects.
+        """
+        if not items:
+            return []
+        if len(items) > 1000:
+            raise ValueError("Cannot retrieve more than 1000 containers at once.")
+
+        result = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url("/models/containers/byids"),
+                method="POST",
+                body=ItemIDBody(items=items),
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[ContainerResponse].model_validate_json(result.success_response.body)
+        return result.items
+
+    def list(
+        self,
+        space: str | None = None,
+        include_global: bool = False,
+        limit: int = 10,
+    ) -> list[ContainerResponse]:
+        """List containers in CDF Project.
+
+        Args:
+            space: If specified, only containers in this space are returned.
+            include_global: If True, include global containers.
+            limit: Maximum number of containers to return. Max is 1000.
+
+        Returns:
+            List of ContainerResponse objects.
+        """
+        if limit > 1000:
+            raise ValueError("Pagination is not (yet) supported for listing containers. The maximum limit is 1000.")
+        parameters: dict[str, PrimitiveType] = {
+            "includeGlobal": include_global,
+            "limit": limit,
+        }
+        if space is not None:
+            parameters["space"] = space
+        result = self._http_client.request_with_retries(
+            ParametersRequest(
+                endpoint_url=self._config.create_api_url("/models/containers"),
+                method="GET",
+                parameters=parameters,
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[ContainerResponse].model_validate_json(result.success_response.body)
+        return result.items

cognite_neat-0.126.1/cognite/neat/_client/data_classes.py

@@ -0,0 +1,10 @@
+from typing import Generic, TypeVar
+
+from pydantic import BaseModel, Field
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class PagedResponse(BaseModel, Generic[T]):
+    items: list[T]
+    next_cursor: str | None = Field(None, alias="nextCursor")

cognite_neat-0.126.1/cognite/neat/_client/data_model_api.py

@@ -0,0 +1,63 @@
+from cognite.neat._data_model.models.dms import DataModelReference, DataModelResponse
+from cognite.neat._utils.http_client import ItemIDBody, ItemsRequest, ParametersRequest
+from cognite.neat._utils.useful_types import PrimitiveType
+
+from .api import NeatAPI
+from .data_classes import PagedResponse
+
+
+class DataModelsAPI(NeatAPI):
+    def retrieve(
+        self,
+        items: list[DataModelReference],
+    ) -> list[DataModelResponse]:
+        """Retrieve data models by their identifiers.
+
+        Args:
+            items: List of data models references identifying the data models to retrieve.
+        Returns:
+            List of DataModelResponse objects.
+        """
+        if not items:
+            return []
+        if len(items) > 1000:
+            raise ValueError("Cannot retrieve more than 1000 containers at once.")
+
+        result = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url("/models/datamodels/byids"),
+                method="POST",
+                body=ItemIDBody(items=items),
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[DataModelResponse].model_validate_json(result.success_response.body)
+        return result.items
+
+    def list(
+        self,
+        space: str | None = None,
+        all_versions: bool = False,
+        include_global: bool = False,
+        limit: int = 10,
+    ) -> list[DataModelResponse]:
+        """List data models in CDF Project."""
+        if limit > 1000:
+            raise ValueError("Pagination is not (yet) supported for listing data models. The maximum limit is 1000.")
+        parameters: dict[str, PrimitiveType] = {
+            "allVersions": all_versions,
+            "includeGlobal": include_global,
+            "limit": limit,
+        }
+        if space is not None:
+            parameters["space"] = space
+        result = self._http_client.request_with_retries(
+            ParametersRequest(
+                endpoint_url=self._config.create_api_url("/models/datamodels"),
+                method="GET",
+                parameters=parameters,
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[DataModelResponse].model_validate_json(result.success_response.body)
+        return result.items

cognite_neat-0.126.1/cognite/neat/_client/spaces_api.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from cognite.neat._data_model.models.dms import SpaceResponse
+from cognite.neat._data_model.models.dms._references import SpaceReference
+from cognite.neat._utils.http_client import ItemIDBody, ItemsRequest, ParametersRequest
+from cognite.neat._utils.useful_types import PrimitiveType
+
+from .api import NeatAPI
+from .data_classes import PagedResponse
+
+
+class SpacesAPI(NeatAPI):
+    def retrieve(self, spaces: list[SpaceReference]) -> list[SpaceResponse]:
+        """Retrieve spaces by their identifiers.
+
+        Args:
+            spaces: List of space identifiers to retrieve.
+
+        Returns:
+            List of SpaceResponse objects.
+        """
+        if not spaces:
+            return []
+        if len(spaces) > 1000:
+            raise ValueError("Cannot retrieve more than 1000 spaces at once.")
+
+        result = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url("/models/spaces/byids"),
+                method="POST",
+                body=ItemIDBody(items=spaces),
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[SpaceResponse].model_validate_json(result.success_response.body)
+        return result.items
+
+    def list(
+        self,
+        include_global: bool = False,
+        limit: int = 10,
+    ) -> list[SpaceResponse]:
+        """List spaces in CDF Project.
+
+        Args:
+            include_global: If True, include global spaces.
+            limit: Maximum number of spaces to return. Max is 1000.
+
+        Returns:
+            List of SpaceResponse objects.
+        """
+        if limit > 1000:
+            raise ValueError("Pagination is not (yet) supported for listing spaces. The maximum limit is 1000.")
+        parameters: dict[str, PrimitiveType] = {
+            "includeGlobal": include_global,
+            "limit": limit,
+        }
+        result = self._http_client.request_with_retries(
+            ParametersRequest(
+                endpoint_url=self._config.create_api_url("/models/spaces"),
+                method="GET",
+                parameters=parameters,
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[SpaceResponse].model_validate_json(result.success_response.body)
+        return result.items

cognite_neat-0.126.1/cognite/neat/_client/views_api.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+from cognite.neat._data_model.models.dms import ViewReference, ViewResponse
+from cognite.neat._utils.http_client import ItemIDBody, ItemsRequest, ParametersRequest
+from cognite.neat._utils.useful_types import PrimitiveType
+
+from .api import NeatAPI
+from .data_classes import PagedResponse
+
+
+class ViewsAPI(NeatAPI):
+    def retrieve(
+        self,
+        items: list[ViewReference],
+        include_inherited_properties: bool = True,
+    ) -> list[ViewResponse]:
+        """Retrieve views by their identifiers.
+
+        Args:
+            items: List of (space, external_id, version) tuples identifying the views to retrieve.
+            include_inherited_properties: If True, include properties inherited from parent views.
+
+        Returns:
+            List of ViewResponse objects.
+        """
+        if not items:
+            return []
+        if len(items) > 1000:
+            raise ValueError("Cannot retrieve more than 1000 views at once.")
+
+        result = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url("/models/views/byids"),
+                method="POST",
+                body=ItemIDBody(items=items),
+                parameters={"includeInheritedProperties": include_inherited_properties},
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[ViewResponse].model_validate_json(result.success_response.body)
+        return result.items
+
+    def list(
+        self,
+        space: str | None = None,
+        all_versions: bool = False,
+        include_inherited_properties: bool = True,
+        include_global: bool = False,
+        limit: int = 10,
+    ) -> list[ViewResponse]:
+        """List views in CDF Project.
+
+        Args:
+            space: If specified, only views in this space are returned.
+            all_versions: If True, return all versions. If False, only return the latest version.
+            include_inherited_properties: If True, include properties inherited from parent views.
+            include_global: If True, include global views.
+            limit: Maximum number of views to return. Max is 1000.
+
+        Returns:
+            List of ViewResponse objects.
+        """
+        if limit > 1000:
+            raise ValueError("Pagination is not (yet) supported for listing views. The maximum limit is 1000.")
+        parameters: dict[str, PrimitiveType] = {
+            "allVersions": all_versions,
+            "includeInheritedProperties": include_inherited_properties,
+            "includeGlobal": include_global,
+            "limit": limit,
+        }
+        if space is not None:
+            parameters["space"] = space
+        result = self._http_client.request_with_retries(
+            ParametersRequest(
+                endpoint_url=self._config.create_api_url("/models/views"),
+                method="GET",
+                parameters=parameters,
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[ViewResponse].model_validate_json(result.success_response.body)
+        return result.items

cognite_neat-0.126.1/cognite/neat/_data_model/_analysis.py

@@ -0,0 +1,127 @@
+from graphlib import TopologicalSorter
+
+from cognite.neat._data_model.models.dms._references import ContainerReference, ViewReference
+from cognite.neat._data_model.models.dms._schema import RequestSchema
+from cognite.neat._data_model.models.dms._view_property import (
+    EdgeProperty,
+    ReverseDirectRelationProperty,
+    ViewCorePropertyRequest,
+)
+from cognite.neat._data_model.models.dms._views import ViewRequest
+
+
+class DataModelAnalysis:
+    def __init__(
+        self,
+        physical: RequestSchema | None = None,
+    ) -> None:
+        self._physical = physical
+
+    @property
+    def physical(self) -> RequestSchema:
+        if self._physical is None:
+            raise ValueError("Physical Data Model is required for this analysis")
+        return self._physical
+
+    def referenced_views(self, include_connection_end_node_types: bool = False) -> set[ViewReference]:
+        """Get all referenced views in the physical data model."""
+        referenced_views = set()
+
+        for view in self.physical.views:
+            referenced_views.add(view.as_reference())
+            if view.implements:
+                for implement in view.implements:
+                    referenced_views.add(implement)
+
+        if include_connection_end_node_types:
+            referenced_views |= set(self.connection_end_node_types.values())
+
+        return referenced_views
+
+    @property
+    def referenced_containers(self) -> set[ContainerReference]:
+        """Get all referenced containers in the physical data model."""
+        referenced_containers = set()
+
+        for view in self.physical.views:
+            for property_ in view.properties.values():
+                if isinstance(property_, ViewCorePropertyRequest):
+                    referenced_containers.add(property_.container)
+
+        for container in self.physical.containers:
+            referenced_containers.add(container.as_reference())
+
+        return referenced_containers
+
+    @property
+    def view_ancestors(self) -> dict[ViewReference, set[ViewReference]]:
+        """Get a mapping of each view to its ancestors in the physical data model."""
+        implements_by_view = self.implements_by_view
+
+        # Topological sort to ensure that concepts include all ancestors
+        for view in list(TopologicalSorter(implements_by_view).static_order()):
+            if view not in implements_by_view:
+                continue
+            implements_by_view[view] |= {
+                grand_parent
+                for parent in implements_by_view[view]
+                for grand_parent in implements_by_view.get(parent, set())
+            }
+
+        return implements_by_view
+
+    @property
+    def implements_by_view(self) -> dict[ViewReference, set[ViewReference]]:
+        """Get a mapping of each view to the views it implements."""
+        implements_mapping: dict[ViewReference, set[ViewReference]] = {}
+
+        for view in self.physical.views:
+            view_ref = view.as_reference()
+            if view_ref not in implements_mapping:
+                implements_mapping[view_ref] = set()
+            if view.implements:
+                for implement in view.implements:
+                    implements_mapping[view_ref].add(implement)
+
+        return implements_mapping
+
+    def view_by_reference(self, include_inherited_properties: bool = True) -> dict[ViewReference, ViewRequest]:
+        """Get a mapping of view references to their corresponding ViewRequest objects."""
+        view_ancestors = self.view_ancestors
+
+        view_by_reference: dict[ViewReference, ViewRequest] = {
+            view.as_reference(): view.model_copy(deep=True) for view in self.physical.views
+        }
+
+        if include_inherited_properties:
+            for ref, view in view_by_reference.items():
+                for ancestor in view_ancestors.get(ref, set()):
+                    if ancestor_view := view_by_reference.get(ancestor):
+                        if ancestor_view.properties:
+                            view.properties.update(ancestor_view.properties)
+
+        return view_by_reference
+
+    @property
+    def connection_end_node_types(self) -> dict[tuple[ViewReference, str], ViewReference]:
+        """Get a mapping of view references to their corresponding ViewRequest objects."""
+        view_by_reference = self.view_by_reference(include_inherited_properties=False)
+        connection_end_node_types: dict[tuple[ViewReference, str], ViewReference] = {}
+
+        for view_ref, view in view_by_reference.items():
+            if not view.properties:
+                continue
+            for prop_ref, property_ in view.properties.items():
+                # direct relation
+                if isinstance(property_, ViewCorePropertyRequest) and property_.source:
+                    connection_end_node_types[(view_ref, prop_ref)] = property_.source
+
+                # reverse direct relation
+                if isinstance(property_, ReverseDirectRelationProperty) and property_.source:
+                    connection_end_node_types[(view_ref, prop_ref)] = property_.source
+
+                # edge property
+                if isinstance(property_, EdgeProperty) and property_.source:
+                    connection_end_node_types[(view_ref, prop_ref)] = property_.source
+
+        return connection_end_node_types

cognite_neat-0.126.1/cognite/neat/_data_model/_constants.py

@@ -0,0 +1,62 @@
+from ._identifiers import NameSpace
+
+XML_SCHEMA_NAMESPACE = NameSpace("http://www.w3.org/2001/XMLSchema#")
+
+CDF_CDM_SPACE = "cdf_cdm"
+CDF_CDM_VERSION = "v1"
+
+COGNITE_CONCEPTS_MAIN = (
+    "CogniteAsset",
+    "CogniteEquipment",
+    "CogniteActivity",
+    "CogniteTimeSeries",
+    "CogniteFile",
+)
+
+COGNITE_CONCEPTS_INTERFACES = (
+    "CogniteDescribable",
+    "CogniteSourceable",
+    "CogniteSchedulable",
+    "CogniteVisualizable",
+)
+
+COGNITE_CONCEPTS_CONFIGURATIONS = (
+    "CogniteSourceSystem",
+    "CogniteUnit",
+    "CogniteAssetClass",
+    "CogniteAssetType",
+    "CogniteEquipmentType",
+    "CogniteFileCategory",
+)
+COGNITE_CONCEPTS_ANNOTATIONS = (
+    "CogniteAnnotation",
+    "CogniteDiagramAnnotation",
+)
+COGNITE_CONCEPTS_3D = (
+    "CogniteCubeMap",
+    "CogniteCADRevision",
+    "CognitePointCloudVolume",
+    "Cognite360ImageAnnotation",
+    "Cognite3DObject",
+    "Cognite3DRevision",
+    "Cognite360Image",
+    "Cognite360ImageCollection",
+    "Cognite360ImageStation",
+    "CognitePointCloudModel",
+    "Cognite3DTransformation",
+    "Cognite360ImageModel",
+    "Cognite3DModel",
+    "CogniteCADModel",
+    "CognitePointCloudRevision",
+    "CogniteCADNode",
+)
+
+COGNITE_CONCEPTS: tuple[str, ...] = (
+    *COGNITE_CONCEPTS_MAIN,
+    *COGNITE_CONCEPTS_INTERFACES,
+    *COGNITE_CONCEPTS_CONFIGURATIONS,
+    *COGNITE_CONCEPTS_ANNOTATIONS,
+    *COGNITE_CONCEPTS_3D,
+)
+
+COGNITE_SPACES = (CDF_CDM_SPACE,)

cognite_neat-0.126.1/cognite/neat/_data_model/_shared.py

@@ -0,0 +1,46 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+from cognite.neat._client.client import NeatClient
+from cognite.neat._issues import IssueList
+
+
+class OnSuccess(ABC):
+    """Abstract base class for post-activity success handlers."""
+
+    def __init__(self, client: NeatClient | None = None) -> None:
+        self._client = client
+        self._issues: IssueList = IssueList()
+        self._results: list = []
+        self._has_run = False
+
+    @property
+    def issues(self) -> IssueList:
+        """List of issues found during the success handler execution."""
+        if not self._has_run:
+            raise RuntimeError("OnSuccess handler has not been run yet.")
+        return self._issues
+
+    @property
+    def results(self) -> list:
+        """List of results produced during the success handler execution."""
+        if not self._has_run:
+            raise RuntimeError("OnSuccess handler has not been run yet.")
+        return self._results
+
+    @abstractmethod
+    def run(self, data_model: Any) -> None:
+        """Execute the success handler on the data model."""
+        pass
+
+
+class OnSuccessIssuesChecker(OnSuccess):
+    """Abstract base class for post-activity success handlers that check for issues of the data model."""
+
+    ...
+
+
+class OnSuccessResultProducer(OnSuccess):
+    """Abstract base class for post-activity success handlers that produce desired outcomes using the data model."""
+
+    ...