cognite-neat 0.126.0__py3-none-any.whl → 0.126.1__py3-none-any.whl

cognite/neat/_utils/http_client/_data_classes.py

@@ -1,45 +1,66 @@
+import sys
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Sequence
-from dataclasses import dataclass, field
-from typing import Generic, Literal, TypeAlias
+from collections import UserList
+from collections.abc import MutableSequence, Sequence
+from typing import Generic, Literal, TypeAlias, TypeVar
 
 import httpx
+from pydantic import BaseModel, ConfigDict, Field, JsonValue, ValidationError, model_serializer
 
+from cognite.neat._exceptions import CDFAPIException
 from cognite.neat._utils.http_client._tracker import ItemsRequestTracker
-from cognite.neat._utils.useful_types import T_ID, JsonVal
+from cognite.neat._utils.useful_types import PrimaryTypes, ReferenceObject, T_Reference
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
 
 StatusCode: TypeAlias = int
 
 
-@dataclass
-class HTTPMessage:
+class HTTPMessage(BaseModel):
     """Base class for HTTP messages (requests and responses)"""
 
-    def dump(self) -> dict[str, JsonVal]:
-        """Dumps the message to a JSON serializable dictionary.
-
-        Returns:
-            dict[str, JsonVal]: The message as a dictionary.
-        """
-        # We avoid using the asdict function as we know we have a shallow structure,
-        # and this roughly ~10x faster.
-        output = self.__dict__.copy()
-        output["type"] = type(self).__name__
-        return output
-
 
-@dataclass
 class FailedRequestMessage(HTTPMessage):
-    error: str
+    message: str
 
 
-@dataclass
 class ResponseMessage(HTTPMessage):
-    status_code: StatusCode
+    code: StatusCode
+    body: str
+
+
+class SuccessResponse(ResponseMessage): ...
+
+
+class ErrorDetails(BaseModel):
+    """This is the structure of failure responses from CDF APIs"""
+
+    code: StatusCode
+    message: str
+    missing: list[JsonValue] | None = None
+    duplicated: list[JsonValue] | None = None
+    is_auto_retryable: bool | None = Field(None, alias="isAutoRetryable")
+
+    @classmethod
+    def from_response(cls, response: httpx.Response) -> "ErrorDetails":
+        try:
+            return _ErrorResponse.model_validate_json(response.text).error
+        except ValidationError:
+            return cls(code=response.status_code, message=response.text)
+
+
+class _ErrorResponse(BaseModel):
+    error: ErrorDetails
+
+
+class FailedResponse(ResponseMessage):
+    error: ErrorDetails
 
 
-@dataclass
-class RequestMessage(HTTPMessage):
+class RequestMessage(HTTPMessage, ABC):
     """Base class for HTTP request messages"""
 
     endpoint_url: str
@@ -54,12 +75,11 @@ class RequestMessage(HTTPMessage):
         return self.connect_attempt + self.read_attempt + self.status_attempt
 
     @abstractmethod
-    def create_responses(
-        self,
-        response: httpx.Response,
-        response_body: dict[str, JsonVal] | None = None,
-        error_message: str | None = None,
-    ) -> Sequence[HTTPMessage]:
+    def create_success_response(self, response: httpx.Response) -> Sequence[HTTPMessage]:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def create_failure_response(self, response: httpx.Response) -> Sequence[HTTPMessage]:
         raise NotImplementedError()
 
     @abstractmethod
@@ -67,158 +87,120 @@ class RequestMessage(HTTPMessage):
         raise NotImplementedError()
 
 
-@dataclass
-class SuccessResponse(ResponseMessage):
-    body: dict[str, JsonVal] | None = None
+class SimpleRequest(RequestMessage):
+    """Base class for requests with a simple success/fail response structure"""
 
+    def create_success_response(self, response: httpx.Response) -> Sequence[HTTPMessage]:
+        return [SuccessResponse(code=response.status_code, body=response.text)]
 
-@dataclass
-class FailedResponse(ResponseMessage):
-    error: str
-    body: dict[str, JsonVal] | None = None
+    def create_failure_response(self, response: httpx.Response) -> Sequence[HTTPMessage]:
+        return [
+            FailedResponse(code=response.status_code, body=response.text, error=ErrorDetails.from_response(response))
+        ]
 
+    def create_failed_request(self, error_message: str) -> Sequence[HTTPMessage]:
+        return [FailedRequestMessage(message=error_message)]
 
-@dataclass
-class SimpleRequest(RequestMessage):
-    """Base class for requests with a simple success/fail response structure"""
 
-    @classmethod
-    def create_responses(
-        cls,
-        response: httpx.Response,
-        response_body: dict[str, JsonVal] | None = None,
-        error_message: str | None = None,
-    ) -> Sequence[ResponseMessage]:
-        if 200 <= response.status_code < 300 and error_message is None:
-            return [SuccessResponse(status_code=response.status_code, body=response_body)]
-        if error_message is None:
-            error_message = f"Request failed with status code {response.status_code}"
-        return [FailedResponse(status_code=response.status_code, error=error_message, body=response_body)]
+class ParametersRequest(SimpleRequest):
+    """Base class for HTTP request messages with query parameters"""
 
-    @classmethod
-    def create_failed_request(cls, error_message: str) -> Sequence[HTTPMessage]:
-        return [FailedRequestMessage(error=error_message)]
+    parameters: dict[str, PrimaryTypes] | None = None
 
 
-@dataclass
-class BodyRequest(RequestMessage, ABC):
+class BodyRequest(ParametersRequest, ABC):
     """Base class for HTTP request messages with a body"""
 
     @abstractmethod
-    def body(self) -> dict[str, JsonVal]:
+    def data(self) -> str:
         raise NotImplementedError()
 
 
-@dataclass
-class ParamRequest(SimpleRequest):
-    """Base class for HTTP request messages with query parameters"""
-
-    parameters: dict[str, str] | None = None
-
+class SimpleBodyRequest(BodyRequest):
+    body: str
 
-@dataclass
-class SimpleBodyRequest(SimpleRequest, BodyRequest):
-    body_content: dict[str, JsonVal] = field(default_factory=dict)
+    def data(self) -> str:
+        return self.body
 
-    def body(self) -> dict[str, JsonVal]:
-        return self.body_content
 
-
-@dataclass
-class ItemMessage:
+class ItemMessage(BaseModel, Generic[T_Reference], ABC):
     """Base class for message related to a specific item"""
 
-    ...
+    ids: Sequence[T_Reference]
 
 
-@dataclass
-class ItemIDMessage(Generic[T_ID], ItemMessage, ABC):
-    """Base class for message related to a specific item identified by an ID"""
+class SuccessResponseItems(ItemMessage[T_Reference], SuccessResponse): ...
 
-    id: T_ID
 
+class FailedResponseItems(ItemMessage[T_Reference], FailedResponse): ...
 
-@dataclass
-class ItemResponse(ItemIDMessage, ResponseMessage, ABC): ...
 
+class FailedRequestItems(ItemMessage[T_Reference], FailedRequestMessage): ...
 
-@dataclass
-class SuccessItem(ItemResponse):
-    item: JsonVal | None = None
 
+T_BaseModel = TypeVar("T_BaseModel", bound=BaseModel)
 
-@dataclass
-class FailedItem(ItemResponse):
-    error: str
 
+class ItemBody(BaseModel, Generic[T_Reference, T_BaseModel], ABC):
+    items: Sequence[T_BaseModel]
+    extra_args: dict[str, JsonValue] | None = None
 
-@dataclass
-class MissingItem(ItemResponse): ...
+    @model_serializer(mode="plain", return_type=dict)
+    def serialize(self) -> dict[str, JsonValue]:
+        data: dict[str, JsonValue] = {
+            "items": [item.model_dump(exclude_unset=True, by_alias=True) for item in self.items]
+        }
+        if isinstance(self.extra_args, dict):
+            data.update(self.extra_args)
+        return data
 
+    @abstractmethod
+    def as_ids(self) -> list[T_Reference]:
+        """Returns the list of item identifiers for the items in the body."""
+        raise NotImplementedError()
 
-@dataclass
-class UnexpectedItem(ItemResponse):
-    item: JsonVal | None = None
+    def split(self, mid: int) -> tuple[Self, Self]:
+        """Splits the body into two smaller bodies.
 
+        This is useful for retrying requests that fail due to size limits or timeouts.
 
-@dataclass
-class FailedRequestItem(ItemIDMessage, FailedRequestMessage): ...
+        Args:
+            mid: The index at which to split the items.
+        Returns:
+            A tuple containing two new ItemBody instances, each with half of the original items.
 
+        """
 
-@dataclass
-class UnknownRequestItem(ItemMessage, FailedRequestMessage):
-    item: JsonVal | None = None
+        type_ = type(self)
+        return type_(items=self.items[:mid], extra_args=self.extra_args), type_(
+            items=self.items[mid:], extra_args=self.extra_args
+        )
 
 
-@dataclass
-class UnknownResponseItem(ItemMessage, ResponseMessage):
-    error: str
-    item: JsonVal | None = None
+class ItemIDBody(ItemBody[ReferenceObject, ReferenceObject]):
+    def as_ids(self) -> list[ReferenceObject]:
+        return list(self.items)
 
 
-@dataclass
-class ItemsRequest(Generic[T_ID], BodyRequest):
+class ItemsRequest(BodyRequest, Generic[T_Reference, T_BaseModel]):
     """Requests message for endpoints that accept multiple items in a single request.
 
     This class provides functionality to split large requests into smaller ones, handle responses for each item,
     and manage errors effectively.
 
     Attributes:
-        items (list[JsonVal]): The list of items to be sent in the request body.
-        extra_body_fields (dict[str, JsonVal]): Additional fields to include in the request body
-        as_id (Callable[[JsonVal], T_ID] | None): A function to extract the ID from each item. If None,
-            IDs are not used.
+        body (ItemBody): The body of the request containing the items to be processed.
         max_failures_before_abort (int): The maximum number of failed split requests before aborting further splits.
 
     """
 
-    items: list[JsonVal] = field(default_factory=list)
-    extra_body_fields: dict[str, JsonVal] = field(default_factory=dict)
-    as_id: Callable[[JsonVal], T_ID] | None = None
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    body: ItemBody[T_Reference, T_BaseModel]
     max_failures_before_abort: int = 50
-    tracker: ItemsRequestTracker | None = field(default=None, init=False)
+    tracker: ItemsRequestTracker | None = None
 
-    def dump(self) -> dict[str, JsonVal]:
-        """Dumps the message to a JSON serializable dictionary.
-
-        This override removes the 'as_id' attribute as it is not serializable.
-
-        Returns:
-            dict[str, JsonVal]: The message as a dictionary.
-        """
-        output = super().dump()
-        if self.as_id is not None:
-            # We cannot serialize functions
-            del output["as_id"]
-        if self.tracker is not None:
-            # We cannot serialize the tracker
-            del output["tracker"]
-        return output
-
-    def body(self) -> dict[str, JsonVal]:
-        if self.extra_body_fields:
-            return {"items": self.items, **self.extra_body_fields}
-        return {"items": self.items}
+    def data(self) -> str:
+        return self.body.model_dump_json(exclude_unset=True, by_alias=True)
 
     def split(self, status_attempts: int) -> "list[ItemsRequest]":
         """Splits the request into two smaller requests.
@@ -235,155 +217,72 @@ class ItemsRequest(Generic[T_ID], BodyRequest):
             A list containing two new ItemsRequest instances, each with half of the original items.
 
         """
-        mid = len(self.items) // 2
+        mid = len(self.body.items) // 2
         if mid == 0:
             return [self]
         tracker = self.tracker or ItemsRequestTracker(self.max_failures_before_abort)
         tracker.register_failure()
-        first_half = ItemsRequest[T_ID](
-            endpoint_url=self.endpoint_url,
-            method=self.method,
-            items=self.items[:mid],
-            extra_body_fields=self.extra_body_fields,
-            as_id=self.as_id,
-            connect_attempt=self.connect_attempt,
-            read_attempt=self.read_attempt,
-            status_attempt=status_attempts,
-        )
-        first_half.tracker = tracker
-        second_half = ItemsRequest[T_ID](
-            endpoint_url=self.endpoint_url,
-            method=self.method,
-            items=self.items[mid:],
-            extra_body_fields=self.extra_body_fields,
-            as_id=self.as_id,
-            connect_attempt=self.connect_attempt,
-            read_attempt=self.read_attempt,
-            status_attempt=status_attempts,
-        )
-        second_half.tracker = tracker
-        return [first_half, second_half]
-
-    def create_responses(
-        self,
-        response: httpx.Response,
-        response_body: dict[str, JsonVal] | None = None,
-        error_message: str | None = None,
-    ) -> Sequence[HTTPMessage]:
+        messages: list[ItemsRequest] = []
+        for body in self.body.split(mid):
+            item_request = ItemsRequest(
+                endpoint_url=self.endpoint_url,
+                method=self.method,
+                body=body,
+                connect_attempt=self.connect_attempt,
+                read_attempt=self.read_attempt,
+                status_attempt=status_attempts,
+            )
+            item_request.tracker = tracker
+            messages.append(item_request)
+        return messages
+
+    def create_success_response(self, response: httpx.Response) -> Sequence[HTTPMessage]:
+        return [SuccessResponseItems(code=response.status_code, body=response.text, ids=self.body.as_ids())]
+
+    def create_failure_response(self, response: httpx.Response) -> Sequence[HTTPMessage]:
         """Creates response messages based on the HTTP response and the original request.
 
         Args:
             response: The HTTP response received from the server.
-            response_body: The parsed JSON body of the response, if available.
-            error_message: An optional error message to use if the response indicates a failure.
-
         Returns:
             A sequence of HTTPMessage instances representing the outcome for each item in the request.
         """
-        if self.as_id is None:
-            return SimpleBodyRequest.create_responses(response, response_body, error_message)
-        request_items_by_id, errors = self._create_items_by_id()
-        responses: list[HTTPMessage] = list(errors)
-        error_message = error_message or "Unknown error"
-
-        if not self._is_items_response(response_body):
-            return self._handle_non_items_response(responses, response, error_message, request_items_by_id)
-
-        # Process items from response
-        if response_body is not None:
-            self._process_response_items(responses, response, response_body, error_message, request_items_by_id)
-
-        # Handle missing items
-        self._handle_missing_items(responses, response, request_items_by_id)
-
-        return responses
-
-    @staticmethod
-    def _handle_non_items_response(
-        responses: list[HTTPMessage],
-        response: httpx.Response,
-        error_message: str,
-        request_items_by_id: dict[T_ID, JsonVal],
-    ) -> list[HTTPMessage]:
-        """Handles responses that do not contain an 'items' field in the body."""
-        if 200 <= response.status_code < 300:
-            responses.extend(
-                SuccessItem(status_code=response.status_code, id=id_) for id_ in request_items_by_id.keys()
-            )
-        else:
-            responses.extend(
-                FailedItem(status_code=response.status_code, error=error_message, id=id_)
-                for id_ in request_items_by_id.keys()
+        return [
+            FailedResponseItems(
+                code=response.status_code,
+                body=response.text,
+                error=ErrorDetails.from_response(response),
+                ids=self.body.as_ids(),
             )
-        return responses
-
-    def _process_response_items(
-        self,
-        responses: list[HTTPMessage],
-        response: httpx.Response,
-        response_body: dict[str, JsonVal],
-        error_message: str,
-        request_items_by_id: dict[T_ID, JsonVal],
-    ) -> None:
-        """Processes each item in the response body and categorizes them based on their status."""
-        for response_item in response_body["items"]:  # type: ignore[union-attr]
-            try:
-                item_id = self.as_id(response_item)  # type: ignore[misc]
-            except Exception as e:
-                responses.append(
-                    UnknownResponseItem(
-                        status_code=response.status_code, item=response_item, error=f"Error extracting ID: {e!s}"
-                    )
-                )
-                continue
-            request_item = request_items_by_id.pop(item_id, None)
-            if request_item is None:
-                responses.append(UnexpectedItem(status_code=response.status_code, id=item_id, item=response_item))
-            elif 200 <= response.status_code < 300:
-                responses.append(SuccessItem(status_code=response.status_code, id=item_id, item=response_item))
-            else:
-                responses.append(FailedItem(status_code=response.status_code, id=item_id, error=error_message))
-
-    @staticmethod
-    def _handle_missing_items(
-        responses: list[HTTPMessage], response: httpx.Response, request_items_by_id: dict[T_ID, JsonVal]
-    ) -> None:
-        """Handles items that were in the request but not present in the response."""
-        for item_id in request_items_by_id.keys():
-            responses.append(MissingItem(status_code=response.status_code, id=item_id))
+        ]
 
     def create_failed_request(self, error_message: str) -> Sequence[HTTPMessage]:
-        if self.as_id is None:
-            return SimpleBodyRequest.create_failed_request(error_message)
-        items_by_id, errors = self._create_items_by_id()
-        results: list[HTTPMessage] = []
-        results.extend(errors)
-        results.extend(FailedRequestItem(id=item_id, error=error_message) for item_id in items_by_id.keys())
-        return results
-
-    def _create_items_by_id(self) -> tuple[dict[T_ID, JsonVal], list[FailedRequestItem | UnknownRequestItem]]:
-        if self.as_id is None:
-            raise ValueError("as_id function must be provided to create items by ID")
-        items_by_id: dict[T_ID, JsonVal] = {}
-        errors: list[FailedRequestItem | UnknownRequestItem] = []
-        for item in self.items:
-            try:
-                item_id = self.as_id(item)
-            except Exception as e:
-                errors.append(UnknownRequestItem(error=f"Error extracting ID: {e!s}", item=item))
-                continue
-            if item_id in items_by_id:
-                errors.append(FailedRequestItem(id=item_id, error=f"Duplicate item ID: {item_id!r}"))
-            else:
-                items_by_id[item_id] = item
-        return items_by_id, errors
-
-    @staticmethod
-    def _is_items_response(body: dict[str, JsonVal] | None = None) -> bool:
-        if body is None:
-            return False
-        if "items" not in body:
-            return False
-        if not isinstance(body["items"], list):
-            return False
-        return True
+        """Creates failed request messages for each item in the request.
+
+        Args:
+            error_message: The error message to include in the failed request messages.
+
+        Returns:
+            A sequence of HTTPMessage instances representing the failed request for each item.
+        """
+        return [FailedRequestItems(message=error_message, ids=self.body.as_ids())]
+
+
+class APIResponse(UserList, MutableSequence[ResponseMessage | FailedRequestMessage]):
+    def __init__(self, collection: Sequence[ResponseMessage | FailedRequestMessage] | None = None):
+        super().__init__(collection or [])
+
+    def raise_for_status(self) -> None:
+        error_messages = [message for message in self.data if not isinstance(message, SuccessResponse)]
+        if error_messages:
+            raise CDFAPIException(error_messages)
+
+    @property
+    def success_response(self) -> SuccessResponse:
+        success = [msg for msg in self.data if isinstance(msg, SuccessResponse)]
+        if len(success) == 1:
+            return success[0]
+        elif success:
+            raise ValueError("Multiple successful HTTP responses found in the messages.")
+        else:
+            raise ValueError("No successful HTTP response found in the messages.")

cognite/neat/_utils/text.py

@@ -1,3 +1,4 @@
+import re
 from collections.abc import Collection
 from typing import Any
 
@@ -60,3 +61,8 @@ def title_case(s: str) -> str:
         'Hello World'
     """
     return " ".join(word.capitalize() for word in s.replace("_", " ").replace("-", " ").split())
+
+
+def split_on_capitals(text: str) -> list[str]:
+    """Split a string at capital letters."""
+    return re.findall(r"[A-Z][a-z]*", text)

cognite/neat/_utils/useful_types.py

@@ -2,13 +2,34 @@ from collections.abc import Hashable
 from datetime import date, datetime, time, timedelta
 from typing import TypeAlias, TypeVar
 
-JsonVal: TypeAlias = None | str | int | float | bool | dict[str, "JsonVal"] | list["JsonVal"]
+from pydantic import BaseModel
+from pydantic.alias_generators import to_camel
 
+JsonVal: TypeAlias = None | str | int | float | bool | dict[str, "JsonVal"] | list["JsonVal"]
+PrimaryTypes: TypeAlias = str | int | float | bool
 
 T_ID = TypeVar("T_ID", bound=Hashable)
-
 # These are the types that openpyxl supports in cells
 CellValueType: TypeAlias = str | int | float | bool | datetime | date | time | timedelta | None
 
 # The format expected for excel sheets representing a data model
 DataModelTableType: TypeAlias = dict[str, list[dict[str, CellValueType]]]
+PrimitiveType: TypeAlias = str | int | float | bool
+
+
+class BaseModelObject(BaseModel, alias_generator=to_camel, extra="ignore"):
+    """Base class for all object. This includes resources and nested objects."""
+
+    ...
+
+
+T_Item = TypeVar("T_Item", bound=BaseModelObject)
+
+
+class ReferenceObject(BaseModelObject, frozen=True, populate_by_name=True):
+    """Base class for all reference objects - these are identifiers."""
+
+    ...
+
+
+T_Reference = TypeVar("T_Reference", bound=ReferenceObject, covariant=True)

cognite/neat/_version.py

@@ -1,2 +1,2 @@
-__version__ = "0.126.0"
+__version__ = "0.126.1"
 __engine__ = "^2.0.4"

cognite/neat/v0/core/_data_model/importers/_rdf/_shared.py

@@ -67,13 +67,13 @@ def parse_concepts(
         else:
             # Handling implements
             if concepts[concept_id]["implements"] and isinstance(concepts[concept_id]["implements"], list):
-                if res["implements"] not in concepts[concept_id]["implements"]:
+                if res["implements"] and res["implements"] not in concepts[concept_id]["implements"]:
                     concepts[concept_id]["implements"].append(res["implements"])
 
             elif concepts[concept_id]["implements"] and isinstance(concepts[concept_id]["implements"], str):
                 concepts[concept_id]["implements"] = [concepts[concept_id]["implements"]]
 
-                if res["implements"] not in concepts[concept_id]["implements"]:
+                if res["implements"] and res["implements"] not in concepts[concept_id]["implements"]:
                     concepts[concept_id]["implements"].append(res["implements"])
             elif res["implements"]:
                 concepts[concept_id]["implements"] = [res["implements"]]

{cognite_neat-0.126.0.dist-info → cognite_neat-0.126.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognite-neat
-Version: 0.126.0
+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/