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

cognite/neat/config.py

@@ -7,7 +7,7 @@ from pathlib import Path
 from typing import Any, Literal, cast
 
 import yaml
-from pydantic import BaseModel, Field, field_serializer, model_validator
+from pydantic import BaseModel, Field, model_validator
 from yaml import safe_load
 
 from cognite.neat.constants import EXAMPLE_GRAPHS, EXAMPLE_RULES, EXAMPLE_WORKFLOWS
@@ -58,7 +58,8 @@ class Config(BaseModel, arbitrary_types_allowed=True):
     log_level: Literal["ERROR", "WARNING", "INFO", "DEBUG"] = "INFO"
     log_format: str = LOG_FORMAT
     download_workflows_from_cdf: bool = Field(
-        default=False, description="Downloads all workflows from CDF automatically and stores them locally"
+        default=False,
+        description="Downloads all workflows from CDF automatically and stores them locally",
     )
     stop_on_error: bool = False
 
@@ -93,30 +94,43 @@ class Config(BaseModel, arbitrary_types_allowed=True):
                 )
         return data
 
-    @staticmethod
-    @field_serializer("cdf_auth_config", when_used="always", return_type=dict)
-    def backwards_compatible_serialize(cdf_auth_config: EnvironmentVariables) -> dict[str, Any]:
-        output: dict[str, Any] = {}
-        config = cdf_auth_config
-        if config.CDF_PROJECT not in {"Missing", "NOT SET"}:
-            output["project"] = config.CDF_PROJECT
-        if config.CDF_CLUSTER not in {"Missing", "NOT SET"}:
-            output["cluster"] = config.CDF_CLUSTER
-        if config.CDF_URL:
-            output["base_url"] = config.CDF_URL
-        if config.IDP_CLIENT_ID:
-            output["client_id"] = config.IDP_CLIENT_ID
-        if config.IDP_CLIENT_SECRET:
-            output["client_secret"] = config.IDP_CLIENT_SECRET
-        if config.IDP_TOKEN_URL:
-            output["token_url"] = config.IDP_TOKEN_URL
-        if config.IDP_SCOPES:
-            output["scopes"] = config.idp_scopes
-        if config.CDF_TIMEOUT:
-            output["timeout"] = config.CDF_TIMEOUT
-        if config.CDF_MAX_WORKERS:
-            output["max_workers"] = config.CDF_MAX_WORKERS
-        return output
+    def as_legacy_config(
+        self,
+    ) -> dict[str, Any]:
+        config: dict[str, Any] = {}
+
+        config["workflows_store_type"] = self.workflows_store_type
+        config["data_store_path"] = str(self.data_store_path)
+        config["workflows_downloader_filter"] = self.workflow_downloader_filter
+
+        config["cdf_client"] = {}
+        if self.cdf_auth_config.CDF_PROJECT not in {"Missing", "NOT SET"}:
+            config["cdf_client"]["project"] = self.cdf_auth_config.CDF_PROJECT
+        if self.cdf_auth_config.CDF_CLUSTER not in {"Missing", "NOT SET"}:
+            config["cdf_client"]["cluster"] = self.cdf_auth_config.CDF_CLUSTER
+        if self.cdf_auth_config.CDF_URL:
+            config["cdf_client"]["base_url"] = self.cdf_auth_config.CDF_URL
+        if self.cdf_auth_config.IDP_CLIENT_ID:
+            config["cdf_client"]["client_id"] = self.cdf_auth_config.IDP_CLIENT_ID
+        if self.cdf_auth_config.IDP_CLIENT_SECRET:
+            config["cdf_client"]["client_secret"] = self.cdf_auth_config.IDP_CLIENT_SECRET
+        if self.cdf_auth_config.IDP_TOKEN_URL:
+            config["cdf_client"]["token_url"] = self.cdf_auth_config.IDP_TOKEN_URL
+        if self.cdf_auth_config.IDP_SCOPES:
+            config["cdf_client"]["scopes"] = self.cdf_auth_config.idp_scopes
+        if self.cdf_auth_config.CDF_TIMEOUT:
+            config["cdf_client"]["timeout"] = self.cdf_auth_config.CDF_TIMEOUT
+        if self.cdf_auth_config.CDF_MAX_WORKERS:
+            config["cdf_client"]["max_workers"] = self.cdf_auth_config.CDF_MAX_WORKERS
+
+        config["cdf_default_dataset_id"] = self.cdf_default_dataset_id
+        config["load_examples"] = self.load_examples
+        config["log_level"] = self.log_level
+        config["log_format"] = self.log_format
+        config["download_workflows_from_cdf"] = self.download_workflows_from_cdf
+        config["stop_on_error"] = self.stop_on_error
+
+        return config
 
     @property
     def _dir_suffix(self) -> str:
@@ -191,7 +205,10 @@ class Config(BaseModel, arbitrary_types_allowed=True):
             ),
             data_store_path=Path(os.environ.get("NEAT_DATA_PATH", "/app/data")),
             cdf_default_dataset_id=int(os.environ.get("NEAT_CDF_DEFAULT_DATASET_ID", 6476640149881990)),
-            log_level=cast(Literal["ERROR", "WARNING", "INFO", "DEBUG"], os.environ.get("NEAT_LOG_LEVEL", "INFO")),
+            log_level=cast(
+                Literal["ERROR", "WARNING", "INFO", "DEBUG"],
+                os.environ.get("NEAT_LOG_LEVEL", "INFO"),
+            ),
             workflow_downloader_filter=workflow_downloader_filter,
             load_examples=bool(os.environ.get("NEAT_LOAD_EXAMPLES", True) in ["True", "true", "1"]),
         )

cognite/neat/exceptions.py

@@ -102,6 +102,12 @@ class InvalidWorkFlowError(NeatException):
         return self.message
 
 
+class NeatValueError(NeatException, ValueError): ...
+
+
+class NeatTypeError(NeatException, TypeError): ...
+
+
 def wrangle_warnings(list_of_warnings: list[WarningMessage]) -> list[dict]:
     warning_list: list[dict] = []
     for warning in list_of_warnings:
@@ -115,12 +121,12 @@ def wrangle_warnings(list_of_warnings: list[WarningMessage]) -> list[dict]:
 def _neat_warning_to_dict(warning: WarningMessage) -> dict:
     category: Any = warning.category
     return {
-        "type": category.type_,
+        "type": category.resource_type,
         "loc": (),
         "msg": str(warning.message),
         "input": None,
         "ctx": dict(
-            type_=category.type_,
+            type_=category.resource_type,
             code=category.code,
             description=category.description,
             example=category.example,

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

@@ -1,6 +1,4 @@
-import json
-import re
-from collections.abc import Callable, Iterable
+from collections.abc import Callable, Iterable, Set
 from datetime import datetime, timezone
 from pathlib import Path
 from typing import cast
@@ -9,17 +7,17 @@ from cognite.client import CogniteClient
 from cognite.client.data_classes import Asset, AssetFilter, AssetList
 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, string_to_ideal_type
+from cognite.neat.utils.auxiliary import create_sha256_hash
 
+from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFExtractor
 
-class AssetsExtractor(BaseExtractor):
+
+class AssetsExtractor(ClassicCDFExtractor[Asset]):
     """Extract data from Cognite Data Fusions Assets into Neat.
 
     Args:
-        assets (Iterable[Asset]): An iterable of assets.
+        items (Iterable[Asset]): An iterable of assets.
         namespace (Namespace, optional): The namespace to use. Defaults to DEFAULT_NAMESPACE.
         to_type (Callable[[Asset], str | None], optional): A function to convert an asset to a type. Defaults to None.
             If None or if the function returns None, the asset will be set to the default type "Asset".
@@ -34,25 +32,7 @@ class AssetsExtractor(BaseExtractor):
             metadata. Defaults to frozenset({"nan", "null", "none", ""}).
     """
 
-    _SPACE_PATTERN = re.compile(r"\s+")
-
-    def __init__(
-        self,
-        assets: Iterable[Asset],
-        namespace: Namespace | None = None,
-        to_type: Callable[[Asset], str | None] | None = None,
-        total: int | None = None,
-        limit: int | None = None,
-        unpack_metadata: bool = True,
-        skip_metadata_values: set[str] | frozenset[str] | None = frozenset({"nan", "null", "none", ""}),
-    ):
-        self.namespace = namespace or DEFAULT_NAMESPACE
-        self.assets = assets
-        self.to_type = to_type
-        self.total = total
-        self.limit = min(limit, total) if limit and total else limit
-        self.unpack_metadata = unpack_metadata
-        self.skip_metadata_values = skip_metadata_values
+    _default_rdf_type = "Asset"
 
     @classmethod
     def from_dataset(
@@ -63,19 +43,18 @@ class AssetsExtractor(BaseExtractor):
         to_type: Callable[[Asset], str | None] | None = None,
         limit: int | None = None,
         unpack_metadata: bool = True,
+        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
     ):
         total = client.assets.aggregate_count(filter=AssetFilter(data_set_ids=[{"externalId": data_set_external_id}]))
 
         return cls(
-            cast(
-                Iterable[Asset],
-                client.assets(data_set_external_ids=data_set_external_id),
-            ),
+            client.assets(data_set_external_ids=data_set_external_id),
             namespace,
             to_type,
             total,
             limit,
             unpack_metadata=unpack_metadata,
+            skip_metadata_values=skip_metadata_values,
         )
 
     @classmethod
@@ -87,6 +66,7 @@ class AssetsExtractor(BaseExtractor):
         to_type: Callable[[Asset], str | None] | None = None,
         limit: int | None = None,
         unpack_metadata: bool = True,
+        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
     ):
         total = client.assets.aggregate_count(
             filter=AssetFilter(asset_subtree_ids=[{"externalId": root_asset_external_id}])
@@ -102,6 +82,7 @@ class AssetsExtractor(BaseExtractor):
             total,
             limit,
             unpack_metadata=unpack_metadata,
+            skip_metadata_values=skip_metadata_values,
         )
 
     @classmethod
@@ -112,44 +93,24 @@ class AssetsExtractor(BaseExtractor):
         to_type: Callable[[Asset], str] | None = None,
         limit: int | None = None,
         unpack_metadata: bool = True,
+        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
     ):
+        assets = AssetList.load(Path(file_path).read_text())
         return cls(
-            AssetList.load(Path(file_path).read_text()),
+            assets,
             namespace,
             to_type,
-            limit,
+            total=len(assets),
+            limit=limit,
             unpack_metadata=unpack_metadata,
+            skip_metadata_values=skip_metadata_values,
         )
 
-    def extract(self) -> Iterable[Triple]:
-        """Extracts an asset with the given asset_id."""
-        if self.total:
-            try:
-                from rich.progress import track
-            except ModuleNotFoundError:
-                to_iterate = self.assets
-            else:
-                to_iterate = track(
-                    self.assets,
-                    total=self.limit or self.total,
-                    description="Extracting Assets",
-                )
-        else:
-            to_iterate = self.assets
-        for no, asset in enumerate(to_iterate):
-            yield from self._asset2triples(asset)
-            if self.limit and no >= self.limit:
-                break
-
-    def _asset2triples(self, asset: Asset) -> list[Triple]:
+    def _item2triples(self, asset: Asset) -> list[Triple]:
         """Converts an asset to triples."""
         id_ = self.namespace[f"Asset_{asset.id}"]
 
-        # Set rdf type
-        type_ = "Asset"
-        if self.to_type:
-            type_ = self.to_type(asset) or type_
-        type_ = self._SPACE_PATTERN.sub("_", type_)
+        type_ = self._get_rdf_type(asset)
 
         triples: list[Triple] = [(id_, RDF.type, self.namespace[type_])]
 
@@ -195,20 +156,7 @@ class AssetsExtractor(BaseExtractor):
                 )
 
         if asset.metadata:
-            if self.unpack_metadata:
-                for key, value in asset.metadata.items():
-                    if value and (
-                        self.skip_metadata_values is None or value.casefold() not in self.skip_metadata_values
-                    ):
-                        triples.append(
-                            (
-                                id_,
-                                self.namespace[key],
-                                Literal(string_to_ideal_type(value)),
-                            )
-                        )
-            else:
-                triples.append((id_, self.namespace.metadata, Literal(json.dumps(asset.metadata))))
+            triples.extend(self._metadata_to_triples(id_, asset.metadata))
 
         # Create connections:
         if asset.parent_id:

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

@@ -0,0 +1,102 @@
+import json
+import re
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Iterable, Set
+from typing import Generic, TypeVar
+
+from cognite.client.data_classes._base import CogniteResource
+from rdflib import Literal, Namespace, URIRef
+
+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
+
+T_CogniteResource = TypeVar("T_CogniteResource", bound=CogniteResource)
+
+DEFAULT_SKIP_METADATA_VALUES = frozenset({"nan", "null", "none", ""})
+
+
+class ClassicCDFExtractor(BaseExtractor, ABC, Generic[T_CogniteResource]):
+    """This is the Base Extractor for all classic CDF resources.
+
+    A classic resource is recognized in that it has a metadata attribute of type dict[str, str].
+
+    Args:
+        items (Iterable[T_CogniteResource]): An iterable of classic resource.
+        namespace (Namespace, optional): The namespace to use. Defaults to DEFAULT_NAMESPACE.
+        to_type (Callable[[T_CogniteResource], 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.
+    """
+
+    _default_rdf_type: str
+    _SPACE_PATTERN = re.compile(r"\s+")
+
+    def __init__(
+        self,
+        items: Iterable[T_CogniteResource],
+        namespace: Namespace | None = None,
+        to_type: Callable[[T_CogniteResource], str | None] | None = None,
+        total: int | None = None,
+        limit: int | None = None,
+        unpack_metadata: bool = True,
+        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
+    ):
+        self.namespace = namespace or DEFAULT_NAMESPACE
+        self.items = items
+        self.to_type = to_type
+        self.total = total
+        self.limit = min(limit, total) if limit and total else limit
+        self.unpack_metadata = unpack_metadata
+        self.skip_metadata_values = skip_metadata_values
+
+    def extract(self) -> Iterable[Triple]:
+        """Extracts an asset with the given asset_id."""
+        if self.total:
+            try:
+                from rich.progress import track
+            except ModuleNotFoundError:
+                to_iterate = self.items
+            else:
+                to_iterate = track(
+                    self.items,
+                    total=self.limit or self.total,
+                    description=f"Extracting {type(self).__name__.removesuffix('Extractor')}",
+                )
+        else:
+            to_iterate = self.items
+        for no, asset in enumerate(to_iterate):
+            yield from self._item2triples(asset)
+            if self.limit and no >= self.limit:
+                break
+
+    @abstractmethod
+    def _item2triples(self, item: T_CogniteResource) -> list[Triple]:
+        raise NotImplementedError()
+
+    def _metadata_to_triples(self, id_: URIRef, metadata: dict[str, str]) -> Iterable[Triple]:
+        if self.unpack_metadata:
+            for key, value in metadata.items():
+                if value and (self.skip_metadata_values is None or value.casefold() not in self.skip_metadata_values):
+                    yield (
+                        id_,
+                        self.namespace[key],
+                        Literal(string_to_ideal_type(value)),
+                    )
+        else:
+            yield id_, self.namespace.metadata, Literal(json.dumps(metadata))
+
+    def _get_rdf_type(self, item: T_CogniteResource) -> str:
+        type_ = self._default_rdf_type
+        if self.to_type:
+            type_ = self.to_type(item) or type_
+        return self._SPACE_PATTERN.sub("_", type_)

cognite/neat/graph/extractors/_classic_cdf/_events.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 Event, EventList
-from pydantic import AnyHttpUrl, ValidationError
-from rdflib import RDF, Literal, Namespace, URIRef
+from cognite.client.data_classes import Event, EventFilter, EventList
+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 EventsExtractor(BaseExtractor):
+
+class EventsExtractor(ClassicCDFExtractor[Event]):
     """Extract data from Cognite Data Fusions Events into Neat.
 
     Args:
-        events (Iterable[Event]): An iterable of events.
+        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.
     """
 
-    def __init__(
-        self,
-        events: Iterable[Event],
-        namespace: Namespace | None = None,
-        unpack_metadata: bool = True,
-    ):
-        self.namespace = namespace or DEFAULT_NAMESPACE
-        self.events = events
-        self.unpack_metadata = unpack_metadata
+    _default_rdf_type = "Event"
 
     @classmethod
     def from_dataset(
@@ -41,15 +38,21 @@ class EventsExtractor(BaseExtractor):
         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,
     ):
+        total = client.events.aggregate_count(filter=EventFilter(data_set_ids=[{"externalId": data_set_external_id}]))
+
         return cls(
-            cast(
-                Iterable[Event],
-                client.events(data_set_external_ids=data_set_external_id),
-            ),
+            client.events(data_set_external_ids=data_set_external_id),
             namespace,
-            unpack_metadata,
+            to_type,
+            total=total,
+            limit=limit,
+            unpack_metadata=unpack_metadata,
+            skip_metadata_values=skip_metadata_values,
         )
 
     @classmethod
@@ -57,20 +60,30 @@ class EventsExtractor(BaseExtractor):
         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,
     ):
-        return cls(EventList.load(Path(file_path).read_text()), namespace, unpack_metadata)
+        events = EventList.load(Path(file_path).read_text())
 
-    def extract(self) -> Iterable[Triple]:
-        """Extract events as triples."""
-        for event in self.events:
-            yield from self._event2triples(event)
+        return cls(
+            events,
+            namespace,
+            to_type,
+            total=len(events),
+            limit=limit,
+            unpack_metadata=unpack_metadata,
+            skip_metadata_values=skip_metadata_values,
+        )
 
-    def _event2triples(self, event: Event) -> list[Triple]:
+    def _item2triples(self, event: Event) -> list[Triple]:
         id_ = self.namespace[f"Event_{event.id}"]
 
+        type_ = self._get_rdf_type(event)
+
         # Set rdf type
-        triples: list[Triple] = [(id_, RDF.type, self.namespace.Event)]
+        triples: list[Triple] = [(id_, RDF.type, self.namespace[type_])]
 
         # Create attributes
 
@@ -87,16 +100,7 @@ class EventsExtractor(BaseExtractor):
             triples.append((id_, self.namespace.subtype, Literal(event.subtype)))
 
         if event.metadata:
-            if self.unpack_metadata:
-                for key, value in event.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(event.metadata))))
+            triples.extend(self._metadata_to_triples(id_, event.metadata))
 
         if event.description:
             triples.append((id_, self.namespace.description, Literal(event.description)))