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

cognite/neat/_graph/extractors/_classic_cdf/_events.py

@@ -1,182 +1,33 @@
-from collections.abc import Callable, Set
-from datetime import datetime, timezone
+from collections.abc import Iterable
 from pathlib import Path
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes import Event, EventFilter, EventList
-from rdflib import RDF, Literal, Namespace
 
-from cognite.neat._shared import Triple
-
-from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFBaseExtractor, InstanceIdPrefix
+from ._base import ClassicCDFBaseExtractor, InstanceIdPrefix
 
 
 class EventsExtractor(ClassicCDFBaseExtractor[Event]):
-    """Extract data from Cognite Data Fusions Events into Neat.
-
-    Args:
-        items (Iterable[Event]): An iterable of items.
-        namespace (Namespace, optional): The namespace to use. Defaults to DEFAULT_NAMESPACE.
-        to_type (Callable[[Event], str | None], optional): A function to convert an item to a type.
-            Defaults to None. If None or if the function returns None, the asset will be set to the default type.
-        total (int, optional): The total number of items to load. If passed, you will get a progress bar if rich
-            is installed. Defaults to None.
-        limit (int, optional): The maximal number of items to load. Defaults to None. This is typically used for
-            testing setup of the extractor. For example, if you are extracting 100 000 assets, you might want to
-            limit the extraction to 1000 assets to test the setup.
-        unpack_metadata (bool, optional): Whether to unpack metadata. Defaults to False, which yields the metadata as
-            a JSON string.
-        skip_metadata_values (set[str] | frozenset[str] | None, optional): If you are unpacking metadata, then
-           values in this set will be skipped.
-    """
+    """Extract data from Cognite Data Fusions Events into Neat."""
 
     _default_rdf_type = "Event"
+    _instance_id_prefix = InstanceIdPrefix.event
 
     @classmethod
-    def from_dataset(
-        cls,
-        client: CogniteClient,
-        data_set_external_id: str,
-        namespace: Namespace | None = None,
-        to_type: Callable[[Event], str | None] | None = None,
-        limit: int | None = None,
-        unpack_metadata: bool = True,
-        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
-    ):
+    def _from_dataset(cls, client: CogniteClient, data_set_external_id: str) -> tuple[int | None, Iterable[Event]]:
         total = client.events.aggregate_count(filter=EventFilter(data_set_ids=[{"externalId": data_set_external_id}]))
-
-        return cls(
-            client.events(data_set_external_ids=data_set_external_id),
-            namespace,
-            to_type,
-            total=total,
-            limit=limit,
-            unpack_metadata=unpack_metadata,
-            skip_metadata_values=skip_metadata_values,
-        )
+        items = client.events(data_set_external_ids=data_set_external_id)
+        return total, items
 
     @classmethod
-    def from_hierarchy(
-        cls,
-        client: CogniteClient,
-        root_asset_external_id: str,
-        namespace: Namespace | None = None,
-        to_type: Callable[[Event], str | None] | None = None,
-        limit: int | None = None,
-        unpack_metadata: bool = True,
-        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
-    ):
+    def _from_hierarchy(cls, client: CogniteClient, root_asset_external_id: str) -> tuple[int | None, Iterable[Event]]:
         total = client.events.aggregate_count(
             filter=EventFilter(asset_subtree_ids=[{"externalId": root_asset_external_id}])
         )
-
-        return cls(
-            client.events(asset_subtree_external_ids=[root_asset_external_id]),
-            namespace,
-            to_type,
-            total,
-            limit,
-            unpack_metadata=unpack_metadata,
-            skip_metadata_values=skip_metadata_values,
-        )
+        items = client.events(asset_subtree_external_ids=[root_asset_external_id])
+        return total, items
 
     @classmethod
-    def from_file(
-        cls,
-        file_path: str,
-        namespace: Namespace | None = None,
-        to_type: Callable[[Event], str | None] | None = None,
-        limit: int | None = None,
-        unpack_metadata: bool = True,
-        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
-    ):
-        events = EventList.load(Path(file_path).read_text())
-
-        return cls(
-            events,
-            namespace,
-            to_type,
-            total=len(events),
-            limit=limit,
-            unpack_metadata=unpack_metadata,
-            skip_metadata_values=skip_metadata_values,
-        )
-
-    def _item2triples(self, event: Event) -> list[Triple]:
-        id_ = self.namespace[f"{InstanceIdPrefix.event}{event.id}"]
-
-        type_ = self._get_rdf_type(event)
-
-        # Set rdf type
-        triples: list[Triple] = [(id_, RDF.type, self.namespace[type_])]
-
-        # Create attributes
-
-        if event.external_id:
-            triples.append((id_, self.namespace.external_id, Literal(event.external_id)))
-
-        if event.source:
-            triples.append((id_, self.namespace.type, Literal(event.source)))
-
-        if event.type:
-            triples.append((id_, self.namespace.type, Literal(event.type)))
-
-        if event.subtype:
-            triples.append((id_, self.namespace.subtype, Literal(event.subtype)))
-
-        if event.metadata:
-            triples.extend(self._metadata_to_triples(id_, event.metadata))
-
-        if event.description:
-            triples.append((id_, self.namespace.description, Literal(event.description)))
-
-        if event.created_time:
-            triples.append(
-                (
-                    id_,
-                    self.namespace.created_time,
-                    Literal(datetime.fromtimestamp(event.created_time / 1000, timezone.utc)),
-                )
-            )
-
-        if event.last_updated_time:
-            triples.append(
-                (
-                    id_,
-                    self.namespace.last_updated_time,
-                    Literal(datetime.fromtimestamp(event.last_updated_time / 1000, timezone.utc)),
-                )
-            )
-
-        if event.start_time:
-            triples.append(
-                (
-                    id_,
-                    self.namespace.start_time,
-                    Literal(datetime.fromtimestamp(event.start_time / 1000, timezone.utc)),
-                )
-            )
-
-        if event.end_time:
-            triples.append(
-                (
-                    id_,
-                    self.namespace.end_time,
-                    Literal(datetime.fromtimestamp(event.end_time / 1000, timezone.utc)),
-                )
-            )
-
-        if event.data_set_id:
-            triples.append(
-                (
-                    id_,
-                    self.namespace.data_set_id,
-                    self.namespace[f"{InstanceIdPrefix.data_set}{event.data_set_id}"],
-                )
-            )
-
-        if event.asset_ids:
-            for asset_id in event.asset_ids:
-                triples.append((id_, self.namespace.asset, self.namespace[f"{InstanceIdPrefix.asset}{asset_id}"]))
-
-        return triples
+    def _from_file(cls, file_path: str | Path) -> tuple[int | None, Iterable[Event]]:
+        assets = EventList.load(Path(file_path).read_text())
+        return len(assets), assets

cognite/neat/_graph/extractors/_classic_cdf/_files.py

@@ -1,202 +1,38 @@
-from collections.abc import Callable, Set
-from datetime import datetime, timezone
+from collections.abc import Iterable
 from pathlib import Path
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes import FileMetadata, FileMetadataFilter, FileMetadataList
-from rdflib import RDF, Literal, Namespace
 
-from cognite.neat._shared import Triple
-
-from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFBaseExtractor, InstanceIdPrefix
-from ._labels import LabelsExtractor
+from ._base import ClassicCDFBaseExtractor, InstanceIdPrefix
 
 
 class FilesExtractor(ClassicCDFBaseExtractor[FileMetadata]):
-    """Extract data from Cognite Data Fusions files metadata into Neat.
-
-    Args:
-        items (Iterable[FileMetadata]): An iterable of items.
-        namespace (Namespace, optional): The namespace to use. Defaults to DEFAULT_NAMESPACE.
-        to_type (Callable[[FileMetadata], str | None], optional): A function to convert an item to a type.
-            Defaults to None. If None or if the function returns None, the asset will be set to the default type.
-        total (int, optional): The total number of items to load. If passed, you will get a progress bar if rich
-            is installed. Defaults to None.
-        limit (int, optional): The maximal number of items to load. Defaults to None. This is typically used for
-            testing setup of the extractor. For example, if you are extracting 100 000 assets, you might want to
-            limit the extraction to 1000 assets to test the setup.
-        unpack_metadata (bool, optional): Whether to unpack metadata. Defaults to False, which yields the metadata as
-            a JSON string.
-        skip_metadata_values (set[str] | frozenset[str] | None, optional): If you are unpacking metadata, then
-           values in this set will be skipped.
-    """
+    """Extract data from Cognite Data Fusions files metadata into Neat."""
 
     _default_rdf_type = "File"
+    _instance_id_prefix = InstanceIdPrefix.file
 
     @classmethod
-    def from_dataset(
+    def _from_dataset(
         cls,
         client: CogniteClient,
         data_set_external_id: str,
-        namespace: Namespace | None = None,
-        to_type: Callable[[FileMetadata], str | None] | None = None,
-        limit: int | None = None,
-        unpack_metadata: bool = True,
-        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
-    ):
-        return cls(
-            client.files(data_set_external_ids=data_set_external_id),
-            namespace=namespace,
-            to_type=to_type,
-            limit=limit,
-            unpack_metadata=unpack_metadata,
-            skip_metadata_values=skip_metadata_values,
-        )
+    ) -> tuple[int | None, Iterable[FileMetadata]]:
+        items = client.files(data_set_external_ids=data_set_external_id)
+        return None, items
 
     @classmethod
-    def from_hierarchy(
-        cls,
-        client: CogniteClient,
-        root_asset_external_id: str,
-        namespace: Namespace | None = None,
-        to_type: Callable[[FileMetadata], str | None] | None = None,
-        limit: int | None = None,
-        unpack_metadata: bool = True,
-        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
-    ):
+    def _from_hierarchy(
+        cls, client: CogniteClient, root_asset_external_id: str
+    ) -> tuple[int | None, Iterable[FileMetadata]]:
         total = client.files.aggregate(
             filter=FileMetadataFilter(asset_subtree_ids=[{"externalId": root_asset_external_id}])
         )[0].count
-
-        return cls(
-            client.files(asset_subtree_external_ids=[root_asset_external_id]),
-            namespace,
-            to_type,
-            total,
-            limit,
-            unpack_metadata=unpack_metadata,
-            skip_metadata_values=skip_metadata_values,
-        )
+        items = client.files(asset_subtree_external_ids=root_asset_external_id)
+        return total, items
 
     @classmethod
-    def from_file(
-        cls,
-        file_path: str,
-        namespace: Namespace | None = None,
-        to_type: Callable[[FileMetadata], str | None] | None = None,
-        limit: int | None = None,
-        unpack_metadata: bool = True,
-        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
-    ):
+    def _from_file(cls, file_path: str | Path) -> tuple[int | None, Iterable[FileMetadata]]:
         file_metadata = FileMetadataList.load(Path(file_path).read_text())
-        return cls(
-            file_metadata,
-            namespace=namespace,
-            to_type=to_type,
-            limit=limit,
-            total=len(file_metadata),
-            unpack_metadata=unpack_metadata,
-            skip_metadata_values=skip_metadata_values,
-        )
-
-    def _item2triples(self, file: FileMetadata) -> list[Triple]:
-        id_ = self.namespace[f"{InstanceIdPrefix.file}{file.id}"]
-
-        type_ = self._get_rdf_type(file)
-
-        # Set rdf type
-        triples: list[Triple] = [(id_, RDF.type, self.namespace[type_])]
-
-        # Create attributes
-
-        if file.external_id:
-            triples.append((id_, self.namespace.external_id, Literal(file.external_id)))
-
-        if file.source:
-            triples.append((id_, self.namespace.type, Literal(file.source)))
-
-        if file.mime_type:
-            triples.append((id_, self.namespace.mime_type, Literal(file.mime_type)))
-
-        if file.uploaded:
-            triples.append((id_, self.namespace.uploaded, Literal(file.uploaded)))
-
-        if file.source:
-            triples.append((id_, self.namespace.source, Literal(file.source)))
-
-        if file.metadata:
-            triples.extend(self._metadata_to_triples(id_, file.metadata))
-
-        if file.source_created_time:
-            triples.append(
-                (
-                    id_,
-                    self.namespace.source_created_time,
-                    Literal(datetime.fromtimestamp(file.source_created_time / 1000, timezone.utc)),
-                )
-            )
-        if file.source_modified_time:
-            triples.append(
-                (
-                    id_,
-                    self.namespace.source_created_time,
-                    Literal(datetime.fromtimestamp(file.source_modified_time / 1000, timezone.utc)),
-                )
-            )
-        if file.uploaded_time:
-            triples.append(
-                (
-                    id_,
-                    self.namespace.uploaded_time,
-                    Literal(datetime.fromtimestamp(file.uploaded_time / 1000, timezone.utc)),
-                )
-            )
-
-        if file.created_time:
-            triples.append(
-                (
-                    id_,
-                    self.namespace.created_time,
-                    Literal(datetime.fromtimestamp(file.created_time / 1000, timezone.utc)),
-                )
-            )
-
-        if file.last_updated_time:
-            triples.append(
-                (
-                    id_,
-                    self.namespace.last_updated_time,
-                    Literal(datetime.fromtimestamp(file.last_updated_time / 1000, timezone.utc)),
-                )
-            )
-
-        if file.labels:
-            for label in file.labels:
-                # external_id can create ill-formed URIs, so we create websafe URIs
-                # since labels do not have internal ids, we use the external_id as the id
-                triples.append(
-                    (
-                        id_,
-                        self.namespace.label,
-                        self.namespace[f"{InstanceIdPrefix.label}{LabelsExtractor._label_id(label)}"],
-                    )
-                )
-
-        if file.security_categories:
-            for category in file.security_categories:
-                triples.append((id_, self.namespace.security_categories, Literal(category)))
-
-        if file.data_set_id:
-            triples.append(
-                (
-                    id_,
-                    self.namespace.data_set_id,
-                    self.namespace[f"{InstanceIdPrefix.data_set}{file.data_set_id}"],
-                )
-            )
-
-        if file.asset_ids:
-            for asset_id in file.asset_ids:
-                triples.append((id_, self.namespace.asset, self.namespace[f"{InstanceIdPrefix.asset}{asset_id}"]))
-
-        return triples
+        return len(file_metadata), file_metadata

cognite/neat/_graph/extractors/_classic_cdf/_labels.py

@@ -1,122 +1,54 @@
-from collections.abc import Callable, Set
-from datetime import datetime, timezone
+from collections.abc import Iterable
 from pathlib import Path
 from urllib.parse import quote
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes import Label, LabelDefinition, LabelDefinitionList
-from rdflib import RDF, Literal, Namespace
 
-from cognite.neat._shared import Triple
-
-from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFBaseExtractor, InstanceIdPrefix
+from ._base import ClassicCDFBaseExtractor, InstanceIdPrefix, T_CogniteResource
 
 
 class LabelsExtractor(ClassicCDFBaseExtractor[LabelDefinition]):
-    """Extract data from Cognite Data Fusions Labels into Neat.
-
-    Args:
-        items (Iterable[LabelDefinition]): An iterable of items.
-        namespace (Namespace, optional): The namespace to use. Defaults to DEFAULT_NAMESPACE.
-        to_type (Callable[[LabelDefinition], str | None], optional): A function to convert an item to a type.
-            Defaults to None. If None or if the function returns None, the asset will be set to the default type.
-        total (int, optional): The total number of items to load. If passed, you will get a progress bar if rich
-            is installed. Defaults to None.
-        limit (int, optional): The maximal number of items to load. Defaults to None. This is typically used for
-            testing setup of the extractor. For example, if you are extracting 100 000 assets, you might want to
-            limit the extraction to 1000 assets to test the setup.
-        unpack_metadata (bool, optional): Whether to unpack metadata. Defaults to False, which yields the metadata as
-            a JSON string.
-        skip_metadata_values (set[str] | frozenset[str] | None, optional): If you are unpacking metadata, then
-           values in this set will be skipped.
-    """
+    """Extract data from Cognite Data Fusions Labels into Neat."""
 
     _default_rdf_type = "Label"
+    _instance_id_prefix = InstanceIdPrefix.label
 
     @classmethod
-    def from_dataset(
-        cls,
-        client: CogniteClient,
-        data_set_external_id: str,
-        namespace: Namespace | None = None,
-        to_type: Callable[[LabelDefinition], str | None] | None = None,
-        limit: int | None = None,
-        unpack_metadata: bool = True,
-        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
-    ):
-        return cls(
-            client.labels(data_set_external_ids=data_set_external_id),
-            namespace=namespace,
-            to_type=to_type,
-            limit=limit,
-            unpack_metadata=unpack_metadata,
-            skip_metadata_values=skip_metadata_values,
-        )
+    def _from_dataset(
+        cls, client: CogniteClient, data_set_external_id: str
+    ) -> tuple[int | None, Iterable[LabelDefinition]]:
+        items = client.labels(data_set_external_ids=data_set_external_id)
+        return None, items
 
     @classmethod
-    def from_file(
-        cls,
-        file_path: str,
-        namespace: Namespace | None = None,
-        to_type: Callable[[LabelDefinition], str | None] | None = None,
-        limit: int | None = None,
-        unpack_metadata: bool = True,
-        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
-    ):
-        labels = LabelDefinitionList.load(Path(file_path).read_text())
-        return cls(
-            labels,
-            total=len(labels),
-            namespace=namespace,
-            to_type=to_type,
-            limit=limit,
-            unpack_metadata=unpack_metadata,
-            skip_metadata_values=skip_metadata_values,
-        )
-
-    def _item2triples(self, label: LabelDefinition) -> list[Triple]:
-        if not label.external_id:
-            return []
-
-        id_ = self.namespace[f"{InstanceIdPrefix.label}{self._label_id(label)}"]
-
-        type_ = self._get_rdf_type(label)
-        # Set rdf type
-        triples: list[Triple] = [(id_, RDF.type, self.namespace[type_])]
+    def _from_hierarchy(
+        cls, client: CogniteClient, root_asset_external_id: str
+    ) -> tuple[int | None, Iterable[T_CogniteResource]]:
+        raise NotImplementedError("Hierarchy is not supported for labels")
 
-        # Create attributes
-        triples.append((id_, self.namespace.external_id, Literal(label.external_id)))
-
-        if label.name:
-            triples.append((id_, self.namespace.name, Literal(label.name)))
-
-        if label.description:
-            triples.append((id_, self.namespace.description, Literal(label.description)))
-
-        if label.created_time:
-            triples.append(
-                (
-                    id_,
-                    self.namespace.created_time,
-                    Literal(datetime.fromtimestamp(label.created_time / 1000, timezone.utc)),
-                )
-            )
-
-        if label.data_set_id:
-            triples.append(
-                (
-                    id_,
-                    self.namespace.data_set_id,
-                    self.namespace[f"{InstanceIdPrefix.data_set}{label.data_set_id}"],
-                )
-            )
+    @classmethod
+    def _from_file(cls, file_path: str | Path) -> tuple[int | None, Iterable[LabelDefinition]]:
+        labels = LabelDefinitionList.load(Path(file_path).read_text())
+        return len(labels), labels
 
-        return triples
+    def _fallback_id(self, item: LabelDefinition) -> str | None:
+        if not item.external_id:
+            return None
+        return self._label_id(item)
 
     @staticmethod
-    def _label_id(label: Label | LabelDefinition) -> str:
+    def _label_id(label: Label | LabelDefinition | dict) -> str:
         # external_id can create ill-formed URIs, so we create websafe URIs
         # since labels do not have internal ids, we use the external_id as the id
-        if label.external_id is None:
+        external_id: str | None = None
+        if isinstance(label, dict):
+            if "externalId" in label:
+                external_id = label["externalId"]
+            elif "external_id" in label:
+                external_id = label["external_id"]
+        else:
+            external_id = label.external_id
+        if external_id is None:
             raise ValueError("External id must be set of the label")
-        return quote(label.external_id)
+        return quote(external_id)