cognite-neat 0.88.0__py3-none-any.whl → 0.88.2__py3-none-any.whl

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

@@ -1,40 +1,37 @@
-import json
-from collections.abc import Iterable
+from collections.abc import Callable, Set
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import cast
 from urllib.parse import quote
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes import FileMetadata, FileMetadataList
-from pydantic import AnyHttpUrl, ValidationError
-from rdflib import RDF, Literal, Namespace, URIRef
+from rdflib import RDF, Literal, Namespace
 
-from cognite.neat.constants import DEFAULT_NAMESPACE
-from cognite.neat.graph.extractors._base import BaseExtractor
 from cognite.neat.graph.models import Triple
-from cognite.neat.utils.auxiliary import string_to_ideal_type
 
+from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFExtractor
 
-class FilesExtractor(BaseExtractor):
+
+class FilesExtractor(ClassicCDFExtractor[FileMetadata]):
     """Extract data from Cognite Data Fusions files metadata into Neat.
 
     Args:
-        files_metadata (Iterable[FileMetadata]): An iterable of files metadata.
+        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.
     """
 
-    def __init__(
-        self,
-        files_metadata: Iterable[FileMetadata],
-        namespace: Namespace | None = None,
-        unpack_metadata: bool = True,
-    ):
-        self.namespace = namespace or DEFAULT_NAMESPACE
-        self.files_metadata = files_metadata
-        self.unpack_metadata = unpack_metadata
+    _default_rdf_type = "File"
 
     @classmethod
     def from_dataset(
@@ -42,15 +39,18 @@ class FilesExtractor(BaseExtractor):
         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(
-            cast(
-                Iterable[FileMetadata],
-                client.files(data_set_external_ids=data_set_external_id),
-            ),
-            namespace,
-            unpack_metadata,
+            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,
         )
 
     @classmethod
@@ -58,24 +58,29 @@ class FilesExtractor(BaseExtractor):
         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,
     ):
+        file_metadata = FileMetadataList.load(Path(file_path).read_text())
         return cls(
-            FileMetadataList.load(Path(file_path).read_text()),
-            namespace,
-            unpack_metadata,
+            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 extract(self) -> Iterable[Triple]:
-        """Extract files metadata as triples."""
-        for event in self.files_metadata:
-            yield from self._file2triples(event)
-
-    def _file2triples(self, file: FileMetadata) -> list[Triple]:
+    def _item2triples(self, file: FileMetadata) -> list[Triple]:
         id_ = self.namespace[f"File_{file.id}"]
 
+        type_ = self._get_rdf_type(file)
+
         # Set rdf type
-        triples: list[Triple] = [(id_, RDF.type, self.namespace.File)]
+        triples: list[Triple] = [(id_, RDF.type, self.namespace[type_])]
 
         # Create attributes
 
@@ -95,16 +100,7 @@ class FilesExtractor(BaseExtractor):
             triples.append((id_, self.namespace.source, Literal(file.source)))
 
         if file.metadata:
-            if self.unpack_metadata:
-                for key, value in file.metadata.items():
-                    if value:
-                        type_aware_value = string_to_ideal_type(value)
-                        try:
-                            triples.append((id_, self.namespace[key], URIRef(str(AnyHttpUrl(type_aware_value)))))  # type: ignore
-                        except ValidationError:
-                            triples.append((id_, self.namespace[key], Literal(type_aware_value)))
-            else:
-                triples.append((id_, self.namespace.metadata, Literal(json.dumps(file.metadata))))
+            triples.extend(self._metadata_to_triples(id_, file.metadata))
 
         if file.source_created_time:
             triples.append(

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

@@ -1,33 +1,37 @@
-from collections.abc import Iterable
+from collections.abc import Callable, Set
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import cast
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes import LabelDefinition, LabelDefinitionList
 from rdflib import RDF, Literal, Namespace
 
-from cognite.neat.constants import DEFAULT_NAMESPACE
-from cognite.neat.graph.extractors._base import BaseExtractor
 from cognite.neat.graph.models import Triple
 from cognite.neat.utils.auxiliary import create_sha256_hash
 
+from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFExtractor
 
-class LabelsExtractor(BaseExtractor):
+
+class LabelsExtractor(ClassicCDFExtractor[LabelDefinition]):
     """Extract data from Cognite Data Fusions Labels into Neat.
 
     Args:
-        labels (Iterable[LabelDefinition]): An iterable of labels.
+        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.
     """
 
-    def __init__(
-        self,
-        labels: Iterable[LabelDefinition],
-        namespace: Namespace | None = None,
-    ):
-        self.namespace = namespace or DEFAULT_NAMESPACE
-        self.labels = labels
+    _default_rdf_type = "Label"
 
     @classmethod
     def from_dataset(
@@ -35,57 +39,76 @@ class LabelsExtractor(BaseExtractor):
         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(
-            cast(
-                Iterable[LabelDefinition],
-                client.labels(data_set_external_ids=data_set_external_id),
-            ),
-            namespace,
+            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,
         )
 
     @classmethod
-    def from_file(cls, file_path: str, namespace: Namespace | None = None):
-        return cls(LabelDefinitionList.load(Path(file_path).read_text()), namespace)
+    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 extract(self) -> Iterable[Triple]:
-        """Extract labels as triples."""
-        for label in self.labels:
-            yield from self._labels2triples(label)
+    def _item2triples(self, label: LabelDefinition) -> list[Triple]:
+        if not label.external_id:
+            return []
 
-    def _labels2triples(self, label: LabelDefinition) -> list[Triple]:
-        if label.external_id:
-            id_ = self.namespace[f"Label_{create_sha256_hash(label.external_id)}"]
+        id_ = self.namespace[f"Label_{create_sha256_hash(label.external_id)}"]
 
-            # Set rdf type
-            triples: list[Triple] = [(id_, RDF.type, self.namespace.Label)]
+        type_ = self._get_rdf_type(label)
+        # Set rdf type
+        triples: list[Triple] = [(id_, RDF.type, self.namespace[type_])]
 
-            # Create attributes
-            triples.append((id_, self.namespace.external_id, Literal(label.external_id)))
+        # 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.name:
+            triples.append((id_, self.namespace.name, Literal(label.name)))
 
-            if label.description:
-                triples.append((id_, self.namespace.description, Literal(label.description)))
+        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.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"Dataset_{label.data_set_id}"],
-                    )
+            )
+
+        if label.data_set_id:
+            triples.append(
+                (
+                    id_,
+                    self.namespace.data_set_id,
+                    self.namespace[f"Dataset_{label.data_set_id}"],
                 )
+            )
 
-            return triples
-        return []
+        return triples

cognite/neat/graph/extractors/_classic_cdf/_relationships.py

@@ -1,34 +1,38 @@
-from collections.abc import Iterable
+from collections.abc import Callable, Set
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import cast
 from urllib.parse import quote
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes import Relationship, RelationshipList
 from rdflib import RDF, Literal, Namespace
 
-from cognite.neat.constants import DEFAULT_NAMESPACE
-from cognite.neat.graph.extractors._base import BaseExtractor
 from cognite.neat.graph.models import Triple
 from cognite.neat.utils.auxiliary import create_sha256_hash
 
+from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFExtractor
 
-class RelationshipsExtractor(BaseExtractor):
+
+class RelationshipsExtractor(ClassicCDFExtractor[Relationship]):
     """Extract data from Cognite Data Fusions Relationships into Neat.
 
     Args:
-        relationships (Iterable[Asset]): An iterable of relationships.
+        items (Iterable[Relationship]): An iterable of items.
         namespace (Namespace, optional): The namespace to use. Defaults to DEFAULT_NAMESPACE.
+        to_type (Callable[[Relationship], 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.
     """
 
-    def __init__(
-        self,
-        relationships: Iterable[Relationship],
-        namespace: Namespace | None = None,
-    ):
-        self.namespace = namespace or DEFAULT_NAMESPACE
-        self.relationships = relationships
+    _default_rdf_type = "Relationship"
 
     @classmethod
     def from_dataset(
@@ -36,33 +40,51 @@ class RelationshipsExtractor(BaseExtractor):
         client: CogniteClient,
         data_set_external_id: str,
         namespace: Namespace | None = None,
+        to_type: Callable[[Relationship], str | None] | None = None,
+        limit: int | None = None,
+        unpack_metadata: bool = True,
+        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
     ):
         return cls(
-            cast(
-                Iterable[Relationship],
-                client.relationships(data_set_external_ids=data_set_external_id),
-            ),
-            namespace,
+            client.relationships(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,
         )
 
     @classmethod
-    def from_file(cls, file_path: str, namespace: Namespace | None = None):
-        return cls(RelationshipList.load(Path(file_path).read_text()), namespace)
-
-    def extract(self) -> Iterable[Triple]:
-        """Extracts an asset with the given asset_id."""
-        for relationship in self.relationships:
-            yield from self._relationship2triples(relationship)
+    def from_file(
+        cls,
+        file_path: str,
+        namespace: Namespace | None = None,
+        to_type: Callable[[Relationship], str | None] | None = None,
+        limit: int | None = None,
+        unpack_metadata: bool = True,
+        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
+    ):
+        relationships = RelationshipList.load(Path(file_path).read_text())
+        return cls(
+            relationships,
+            namespace=namespace,
+            total=len(relationships),
+            to_type=to_type,
+            limit=limit,
+            unpack_metadata=unpack_metadata,
+            skip_metadata_values=skip_metadata_values,
+        )
 
-    def _relationship2triples(self, relationship: Relationship) -> list[Triple]:
+    def _item2triples(self, relationship: Relationship) -> list[Triple]:
         """Converts an asset to triples."""
 
         if relationship.external_id and relationship.source_external_id and relationship.target_external_id:
             # relationships do not have an internal id, so we generate one
             id_ = self.namespace[f"Relationship_{create_sha256_hash(relationship.external_id)}"]
 
+            type_ = self._get_rdf_type(relationship)
             # Set rdf type
-            triples: list[Triple] = [(id_, RDF.type, self.namespace["Relationship"])]
+            triples: list[Triple] = [(id_, RDF.type, self.namespace[type_])]
 
             # Set source and target types
             if source_type := relationship.source_type:

cognite/neat/graph/extractors/_classic_cdf/_sequences.py

@@ -1,39 +1,36 @@
-import json
-from collections.abc import Iterable
+from collections.abc import Callable, Set
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import cast
 
 from cognite.client import CogniteClient
-from cognite.client.data_classes import Sequence, SequenceList
-from pydantic import AnyHttpUrl, ValidationError
-from rdflib import RDF, Literal, Namespace, URIRef
+from cognite.client.data_classes import Sequence, SequenceFilter, SequenceList
+from rdflib import RDF, Literal, Namespace
 
-from cognite.neat.constants import DEFAULT_NAMESPACE
-from cognite.neat.graph.extractors._base import BaseExtractor
 from cognite.neat.graph.models import Triple
-from cognite.neat.utils.auxiliary import string_to_ideal_type
 
+from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFExtractor
 
-class SequencesExtractor(BaseExtractor):
+
+class SequencesExtractor(ClassicCDFExtractor[Sequence]):
     """Extract data from Cognite Data Fusions Sequences into Neat.
 
     Args:
-        sequence (Iterable[Sequence]): An iterable of sequences.
+        items (Iterable[Sequence]): An iterable of items.
         namespace (Namespace, optional): The namespace to use. Defaults to DEFAULT_NAMESPACE.
+        to_type (Callable[[Sequence], 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.
     """
 
-    def __init__(
-        self,
-        sequence: Iterable[Sequence],
-        namespace: Namespace | None = None,
-        unpack_metadata: bool = True,
-    ):
-        self.namespace = namespace or DEFAULT_NAMESPACE
-        self.sequence = sequence
-        self.unpack_metadata = unpack_metadata
+    _default_rdf_type = "Sequence"
 
     @classmethod
     def from_dataset(
@@ -41,15 +38,22 @@ class SequencesExtractor(BaseExtractor):
         client: CogniteClient,
         data_set_external_id: str,
         namespace: Namespace | None = None,
+        to_type: Callable[[Sequence], str | None] | None = None,
+        limit: int | None = None,
         unpack_metadata: bool = True,
+        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
     ):
+        total = client.sequences.aggregate_count(
+            filter=SequenceFilter(data_set_ids=[{"externalId": data_set_external_id}])
+        )
         return cls(
-            cast(
-                Iterable[Sequence],
-                client.sequences(data_set_external_ids=data_set_external_id),
-            ),
-            namespace,
-            unpack_metadata,
+            client.sequences(data_set_external_ids=data_set_external_id),
+            total=total,
+            namespace=namespace,
+            to_type=to_type,
+            limit=limit,
+            unpack_metadata=unpack_metadata,
+            skip_metadata_values=skip_metadata_values,
         )
 
     @classmethod
@@ -57,20 +61,28 @@ class SequencesExtractor(BaseExtractor):
         cls,
         file_path: str,
         namespace: Namespace | None = None,
+        to_type: Callable[[Sequence], str | None] | None = None,
+        limit: int | None = None,
         unpack_metadata: bool = True,
+        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
     ):
-        return cls(SequenceList.load(Path(file_path).read_text()), namespace, unpack_metadata)
-
-    def extract(self) -> Iterable[Triple]:
-        """Extract sequences as triples."""
-        for sequence in self.sequence:
-            yield from self._sequence2triples(sequence)
+        sequences = SequenceList.load(Path(file_path).read_text())
+        return cls(
+            sequences,
+            total=len(sequences),
+            namespace=namespace,
+            to_type=to_type,
+            limit=limit,
+            unpack_metadata=unpack_metadata,
+            skip_metadata_values=skip_metadata_values,
+        )
 
-    def _sequence2triples(self, sequence: Sequence) -> list[Triple]:
+    def _item2triples(self, sequence: Sequence) -> list[Triple]:
         id_ = self.namespace[f"Sequence_{sequence.id}"]
 
+        type_ = self._get_rdf_type(sequence)
         # Set rdf type
-        triples: list[Triple] = [(id_, RDF.type, self.namespace.Sequence)]
+        triples: list[Triple] = [(id_, RDF.type, self.namespace[type_])]
 
         # Create attributes
 
@@ -81,22 +93,7 @@ class SequencesExtractor(BaseExtractor):
             triples.append((id_, self.namespace.name, Literal(sequence.name)))
 
         if sequence.metadata:
-            if self.unpack_metadata:
-                for key, value in sequence.metadata.items():
-                    if value:
-                        type_aware_value = string_to_ideal_type(value)
-                        try:
-                            triples.append((id_, self.namespace[key], URIRef(str(AnyHttpUrl(type_aware_value)))))  # type: ignore
-                        except ValidationError:
-                            triples.append((id_, self.namespace[key], Literal(type_aware_value)))
-            else:
-                triples.append(
-                    (
-                        id_,
-                        self.namespace.metadata,
-                        Literal(json.dumps(sequence.metadata)),
-                    )
-                )
+            triples.extend(self._metadata_to_triples(id_, sequence.metadata))
 
         if sequence.description:
             triples.append((id_, self.namespace.description, Literal(sequence.description)))