cognite-neat 0.84.0__tar.gz → 0.85.0__tar.gz

{cognite_neat-0.84.0 → cognite_neat-0.85.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cognite-neat
-Version: 0.84.0
+Version: 0.85.0
 Summary: Knowledge graph transformation
 Home-page: https://cognite-neat.readthedocs-hosted.com/
 License: Apache-2.0
@@ -16,6 +16,7 @@ Provides-Extra: all
 Provides-Extra: docs
 Provides-Extra: google
 Provides-Extra: graphql
+Provides-Extra: jupyter
 Provides-Extra: oxi
 Provides-Extra: service
 Requires-Dist: PyYAML
@@ -46,6 +47,7 @@ Requires-Dist: pymdown-extensions ; extra == "docs"
 Requires-Dist: pyoxigraph (==0.3.19) ; extra == "oxi" or extra == "all"
 Requires-Dist: rdflib
 Requires-Dist: requests
+Requires-Dist: rich[jupyter] (>=13.7.1,<14.0.0) ; extra == "jupyter"
 Requires-Dist: schedule (>=1,<2) ; extra == "service" or extra == "all"
 Requires-Dist: tomli (>=2.0.1,<3.0.0) ; python_version < "3.11"
 Requires-Dist: typing_extensions (>=4.8,<5.0) ; python_version < "3.11"

cognite_neat-0.85.0/cognite/neat/__init__.py

@@ -0,0 +1,4 @@
+from ._version import __version__
+from .utils.auth import get_cognite_client
+
+__all__ = ["__version__", "get_cognite_client"]

{cognite_neat-0.84.0 → cognite_neat-0.85.0}/cognite/neat/_shared.py

@@ -17,8 +17,24 @@ class NeatObject:
         """Return a dictionary representation of the object."""
         raise NotImplementedError()
 
+    def _repr_html_(self) -> str:
+        return pd.Series(self.dump(aggregate=True)).to_frame(name="value")._repr_html_()
+
+
+@dataclass(frozen=True)
+class FrozenNeatObject:
+    """A frozen neat object can be dumped to a dictionary."""
+
+    @abstractmethod
+    def dump(self, aggregate: bool = True) -> dict[str, Any]:
+        """Return a dictionary representation of the object."""
+        raise NotImplementedError()
+
+    def _repr_html_(self) -> str:
+        return pd.Series(self.dump(aggregate=True)).to_frame(name="value")._repr_html_()
+
 
-T_NeatObject = TypeVar("T_NeatObject", bound=NeatObject)
+T_NeatObject = TypeVar("T_NeatObject", bound=NeatObject | FrozenNeatObject)
 
 
 class NeatList(list, Sequence[T_NeatObject]):

cognite_neat-0.85.0/cognite/neat/_version.py

@@ -0,0 +1 @@
+__version__ = "0.85.0"

{cognite_neat-0.84.0 → cognite_neat-0.85.0}/cognite/neat/graph/extractors/__init__.py

@@ -1,3 +1,4 @@
+from ._base import BaseExtractor
 from ._classic_cdf._assets import AssetsExtractor
 from ._classic_cdf._events import EventsExtractor
 from ._classic_cdf._files import FilesExtractor
@@ -10,6 +11,7 @@ from ._mock_graph_generator import MockGraphGenerator
 from ._rdf_file import RdfFileExtractor
 
 __all__ = [
+    "BaseExtractor",
     "AssetsExtractor",
     "MockGraphGenerator",
     "RelationshipsExtractor",
@@ -35,3 +37,23 @@ TripleExtractors = (
     | RdfFileExtractor
     | DexpiExtractor
 )
+
+
+def _repr_html_() -> str:
+    import pandas as pd
+
+    table = pd.DataFrame(  # type: ignore[operator]
+        [
+            {
+                "Extractor": name,
+                "Description": globals()[name].__doc__.strip().split("\n")[0] if globals()[name].__doc__ else "Missing",
+            }
+            for name in __all__
+            if name != "BaseExtractor"
+        ]
+    )._repr_html_()
+
+    return (
+        "<strong>Extractor</strong> An extractor is used to read data from "
+        f"a source into Neat's internal triple storage. <br />{table}"
+    )

{cognite_neat-0.84.0 → cognite_neat-0.85.0}/cognite/neat/graph/extractors/_base.py

@@ -2,6 +2,7 @@ from abc import abstractmethod
 from collections.abc import Iterable
 
 from cognite.neat.graph.models import Triple
+from cognite.neat.utils.auxiliary import class_html_doc
 
 
 class BaseExtractor:
@@ -12,3 +13,7 @@ class BaseExtractor:
     @abstractmethod
     def extract(self) -> Iterable[Triple]:
         raise NotImplementedError()
+
+    @classmethod
+    def _repr_html_(cls) -> str:
+        return class_html_doc(cls)

{cognite_neat-0.84.0 → cognite_neat-0.85.0}/cognite/neat/graph/extractors/_classic_cdf/_assets.py

@@ -14,6 +14,13 @@ from cognite.neat.utils.utils import create_sha256_hash, string_to_ideal_type
 
 
 class AssetsExtractor(BaseExtractor):
+    """Extract data from Cognite Data Fusions Assets into Neat.
+
+    Args:
+        assets (Iterable[Asset]): An iterable of assets.
+        namespace (Namespace, optional): The namespace to use. Defaults to DEFAULT_NAMESPACE.
+    """
+
     def __init__(
         self,
         assets: Iterable[Asset],

{cognite_neat-0.84.0 → cognite_neat-0.85.0}/cognite/neat/graph/extractors/_classic_cdf/_events.py

@@ -15,6 +15,13 @@ from cognite.neat.utils.utils import string_to_ideal_type
 
 
 class EventsExtractor(BaseExtractor):
+    """Extract data from Cognite Data Fusions Events into Neat.
+
+    Args:
+        events (Iterable[Event]): An iterable of events.
+        namespace (Namespace, optional): The namespace to use. Defaults to DEFAULT_NAMESPACE.
+    """
+
     def __init__(
         self,
         events: Iterable[Event],

{cognite_neat-0.84.0 → cognite_neat-0.85.0}/cognite/neat/graph/extractors/_classic_cdf/_labels.py

@@ -14,6 +14,13 @@ from cognite.neat.utils.utils import create_sha256_hash
 
 
 class LabelsExtractor(BaseExtractor):
+    """Extract data from Cognite Data Fusions Labels into Neat.
+
+    Args:
+        labels (Iterable[LabelDefinition]): An iterable of labels.
+        namespace (Namespace, optional): The namespace to use. Defaults to DEFAULT_NAMESPACE.
+    """
+
     def __init__(
         self,
         labels: Iterable[LabelDefinition],

{cognite_neat-0.84.0 → cognite_neat-0.85.0}/cognite/neat/graph/extractors/_classic_cdf/_relationships.py

@@ -15,6 +15,13 @@ from cognite.neat.utils.utils import create_sha256_hash
 
 
 class RelationshipsExtractor(BaseExtractor):
+    """Extract data from Cognite Data Fusions Relationships into Neat.
+
+    Args:
+        relationships (Iterable[Asset]): An iterable of relationships.
+        namespace (Namespace, optional): The namespace to use. Defaults to DEFAULT_NAMESPACE.
+    """
+
     def __init__(
         self,
         relationships: Iterable[Relationship],

{cognite_neat-0.84.0 → cognite_neat-0.85.0}/cognite/neat/graph/extractors/_classic_cdf/_sequences.py

@@ -15,6 +15,13 @@ from cognite.neat.utils.utils import string_to_ideal_type
 
 
 class SequencesExtractor(BaseExtractor):
+    """Extract data from Cognite Data Fusions Sequences into Neat.
+
+    Args:
+        sequence (Iterable[Sequence]): An iterable of sequences.
+        namespace (Namespace, optional): The namespace to use. Defaults to DEFAULT_NAMESPACE.
+    """
+
     def __init__(
         self,
         sequence: Iterable[Sequence],

{cognite_neat-0.84.0 → cognite_neat-0.85.0}/cognite/neat/graph/extractors/_classic_cdf/_timeseries.py

@@ -15,6 +15,13 @@ from cognite.neat.utils.utils import string_to_ideal_type
 
 
 class TimeSeriesExtractor(BaseExtractor):
+    """Extract data from Cognite Data Fusions TimeSeries into Neat.
+
+    Args:
+        timeseries (Iterable[TimeSeries]): An iterable of timeseries.
+        namespace (Namespace, optional): The namespace to use. Defaults to DEFAULT_NAMESPACE.
+    """
+
     def __init__(
         self,
         timeseries: Iterable[TimeSeries],

{cognite_neat-0.84.0 → cognite_neat-0.85.0}/cognite/neat/graph/extractors/_dexpi.py

@@ -10,7 +10,7 @@ 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.utils import as_neat_compliant_uri
-from cognite.neat.utils.xml import get_children, iterate_tree
+from cognite.neat.utils.xml_ import get_children, iterate_tree
 
 DEXPI = Namespace("http://sandbox.dexpi.org/rdl/")
 QUDT = Namespace("https://qudt.org/vocab/unit/")

{cognite_neat-0.84.0 → cognite_neat-0.85.0}/cognite/neat/graph/extractors/_rdf_file.py

@@ -7,6 +7,14 @@ from cognite.neat.graph.extractors._base import BaseExtractor
 
 
 class RdfFileExtractor(BaseExtractor):
+    """Extract data from RDF files into Neat.
+
+    Args:
+        filepath (Path): The path to the RDF file.
+        mime_type (MIMETypes, optional): The MIME type of the RDF file. Defaults to "application/rdf+xml".
+        base_uri (URIRef, optional): The base URI to use. Defaults to None.
+    """
+
     def __init__(
         self,
         filepath: Path,

cognite_neat-0.85.0/cognite/neat/graph/loaders/__init__.py

@@ -0,0 +1,23 @@
+from ._base import BaseLoader, CDFLoader
+from ._rdf2dms import DMSLoader
+
+__all__ = ["BaseLoader", "CDFLoader", "DMSLoader"]
+
+
+def _repr_html_() -> str:
+    import pandas as pd
+
+    table = pd.DataFrame(  # type: ignore[operator]
+        [
+            {
+                "Loader": name,
+                "Description": globals()[name].__doc__.strip().split("\n")[0] if globals()[name].__doc__ else "Missing",
+            }
+            for name in __all__
+            if name not in ("BaseLoader", "CDFLoader")
+        ]
+    )._repr_html_()
+
+    return (
+        "<strong>Loader</strong> A loader writes data from Neat's triple storage into a target system" f"<br />{table}"
+    )

cognite_neat-0.85.0/cognite/neat/graph/loaders/_base.py

@@ -0,0 +1,86 @@
+from abc import ABC, abstractmethod
+from collections.abc import Hashable, Iterable
+from pathlib import Path
+from typing import ClassVar, Generic, TypeVar
+
+from cognite.client import CogniteClient
+from cognite.client.data_classes.capabilities import Capability
+
+from cognite.neat.graph import NeatGraphStore
+from cognite.neat.graph.issues.loader import FailedAuthorizationError
+from cognite.neat.issues import NeatIssue, NeatIssueList
+from cognite.neat.utils.auxiliary import class_html_doc
+from cognite.neat.utils.upload import UploadResult, UploadResultList
+
+T_Output = TypeVar("T_Output")
+
+
+class BaseLoader(ABC, Generic[T_Output]):
+    _new_line = "\n"
+    _encoding = "utf-8"
+
+    def __init__(self, graph_store: NeatGraphStore):
+        self.graph_store = graph_store
+
+    @abstractmethod
+    def write_to_file(self, filepath: Path) -> None:
+        raise NotImplementedError
+
+    def load(self, stop_on_exception: bool = False) -> Iterable[T_Output | NeatIssue]:
+        """Load the graph with data."""
+        return self._load(stop_on_exception)
+
+    @abstractmethod
+    def _load(self, stop_on_exception: bool = False) -> Iterable[T_Output | NeatIssue]:
+        """Load the graph with data."""
+        pass
+
+    @classmethod
+    def _repr_html_(cls) -> str:
+        return class_html_doc(cls)
+
+
+class CDFLoader(BaseLoader[T_Output]):
+    _UPLOAD_BATCH_SIZE: ClassVar[int] = 1000
+
+    def load_into_cdf(self, client: CogniteClient, dry_run: bool = False) -> UploadResultList:
+        return UploadResultList(self.load_into_cdf_iterable(client, dry_run))
+
+    def load_into_cdf_iterable(self, client: CogniteClient, dry_run: bool = False) -> Iterable[UploadResult]:
+        missing_capabilities = client.iam.verify_capabilities(self._get_required_capabilities())
+        if missing_capabilities:
+            upload_result = UploadResult[Hashable](name=type(self).__name__)
+            upload_result.issues.append(
+                FailedAuthorizationError(action="Upload to CDF", reason=str(missing_capabilities))
+            )
+            yield upload_result
+            return
+
+        issues = NeatIssueList[NeatIssue]()
+        items: list[T_Output] = []
+        for result in self.load(stop_on_exception=False):
+            if isinstance(result, NeatIssue):
+                issues.append(result)
+            else:
+                items.append(result)
+
+            if len(items) >= self._UPLOAD_BATCH_SIZE:
+                yield self._upload_to_cdf(client, items, dry_run, issues)
+                issues = NeatIssueList[NeatIssue]()
+                items = []
+        if items:
+            yield self._upload_to_cdf(client, items, dry_run, issues)
+
+    @abstractmethod
+    def _get_required_capabilities(self) -> list[Capability]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def _upload_to_cdf(
+        self,
+        client: CogniteClient,
+        items: list[T_Output],
+        dry_run: bool,
+        read_issues: NeatIssueList,
+    ) -> UploadResult:
+        raise NotImplementedError

cognite_neat-0.85.0/cognite/neat/graph/loaders/_rdf2asset.py

@@ -0,0 +1,123 @@
+from collections.abc import Sequence
+from dataclasses import dataclass, fields
+
+from cognite.client import CogniteClient
+from cognite.client.data_classes import AssetWrite
+
+from cognite.neat.graph._tracking.base import Tracker
+from cognite.neat.graph._tracking.log import LogTracker
+from cognite.neat.graph.stores import NeatGraphStore
+from cognite.neat.issues import NeatIssue, NeatIssueList
+from cognite.neat.rules.models import AssetRules
+
+from ._base import CDFLoader
+
+
+@dataclass(frozen=True)
+class AssetLoaderMetadataKeys:
+    """Class holding mapping between NEAT metadata key names and their desired names
+    in CDF Asset metadata
+
+    Args:
+        start_time: Start time key name
+        end_time: End time key name
+        update_time: Update time key name
+        resurrection_time: Resurrection time key name
+        identifier: Identifier key name
+        active: Active key name
+        type: Type key name
+    """
+
+    start_time: str = "start_time"
+    end_time: str = "end_time"
+    update_time: str = "update_time"
+    resurrection_time: str = "resurrection_time"
+    identifier: str = "identifier"
+    active: str = "active"
+    type: str = "type"
+
+    def as_aliases(self) -> dict[str, str]:
+        return {str(field.default): getattr(self, field.name) for field in fields(self)}
+
+
+class AssetLoader(CDFLoader[AssetWrite]):
+    def __init__(
+        self,
+        rules: AssetRules,
+        graph_store: NeatGraphStore,
+        data_set_id: int,
+        use_orphanage: bool = False,
+        use_labels: bool = False,
+        asset_external_id_prefix: str | None = None,
+        metadata_keys: AssetLoaderMetadataKeys | None = None,
+        create_issues: Sequence[NeatIssue] | None = None,
+        tracker: type[Tracker] | None = None,
+    ):
+        super().__init__(graph_store)
+
+        self.rules = rules
+        self.data_set_id = data_set_id
+        self.use_labels = use_labels
+        self.use_orphanage = use_orphanage
+
+        self.orphanage_external_id = (
+            f"{asset_external_id_prefix or ''}orphanage-{data_set_id}" if use_orphanage else None
+        )
+
+        self.asset_external_id_prefix = asset_external_id_prefix
+        self.metadata_keys = metadata_keys or AssetLoaderMetadataKeys()
+
+        self._issues = NeatIssueList[NeatIssue](create_issues or [])
+        self._tracker: type[Tracker] = tracker or LogTracker
+
+    @classmethod
+    def from_rules(
+        cls,
+        rules: AssetRules,
+        graph_store: NeatGraphStore,
+        data_set_id: int,
+        use_orphanage: bool = False,
+        use_labels: bool = False,
+        asset_external_id_prefix: str | None = None,
+        metadata_keys: AssetLoaderMetadataKeys | None = None,
+    ) -> "AssetLoader":
+        issues: list[NeatIssue] = []
+
+        return cls(
+            rules, graph_store, data_set_id, use_orphanage, use_labels, asset_external_id_prefix, metadata_keys, issues
+        )
+
+    def _create_validation_classes(self) -> None:
+        # need to get back class-property pairs where are definition of
+        # asset implementations, extend InformationRulesAnalysis make it generic
+
+        # by default if there is not explicitly stated external_id
+        # use rdf:type and drop the prefix
+
+        # based on those create pydantic model AssetDefinition
+        # which will have .to_asset_write()
+
+        raise NotImplementedError("Not implemented yet, this is placeholder")
+
+    def categorize_assets(self, client: CogniteClient) -> None:
+        """Categorize assets to those to be created, updated, decommissioned, or resurrected"""
+
+        raise NotImplementedError("Not implemented yet, this is placeholder")
+
+    def load_to_cdf(self, client: CogniteClient, dry_run: bool = False) -> Sequence[AssetWrite]:
+        # generate assets
+        # check for circular asset hierarchy
+        # check for orphaned assets
+        # batch upsert of assets to CDF (otherwise we will hit the API rate limit)
+
+        raise NotImplementedError("Not implemented yet, this is placeholder")
+
+    @classmethod
+    def _check_for_circular_asset_hierarchy(cls, assets: list[AssetWrite]) -> None:
+        """Check for circular references in the asset rules"""
+        raise NotImplementedError("Not implemented yet, this is placeholder")
+
+    @classmethod
+    def _check_for_orphaned_assets(cls, assets: list[AssetWrite]) -> None:
+        """Check for circular references in the asset rules"""
+        raise NotImplementedError("Not implemented yet, this is placeholder")

{cognite_neat-0.84.0 → cognite_neat-0.85.0}/cognite/neat/graph/loaders/_rdf2dms.py

@@ -10,6 +10,7 @@ from cognite.client import CogniteClient
 from cognite.client import data_modeling as dm
 from cognite.client.data_classes.capabilities import Capability, DataModelInstancesAcl
 from cognite.client.data_classes.data_modeling import ViewId
+from cognite.client.data_classes.data_modeling.ids import InstanceId
 from cognite.client.data_classes.data_modeling.views import SingleEdgeConnection
 from cognite.client.exceptions import CogniteAPIError
 from pydantic import ValidationInfo, create_model, field_validator
@@ -21,27 +22,38 @@ from cognite.neat.graph.stores import NeatGraphStore
 from cognite.neat.issues import NeatIssue, NeatIssueList
 from cognite.neat.rules.models import DMSRules
 from cognite.neat.rules.models.data_types import _DATA_TYPE_BY_DMS_TYPE
-from cognite.neat.utils.upload import UploadDiffsID
+from cognite.neat.utils.upload import UploadResult
 from cognite.neat.utils.utils import create_sha256_hash
 
 from ._base import CDFLoader
 
 
 class DMSLoader(CDFLoader[dm.InstanceApply]):
+    """Load data from Cognite Data Fusions Data Modeling Service (DMS) into Neat.
+
+    Args:
+        graph_store (NeatGraphStore): The graph store to load the data into.
+        data_model (dm.DataModel[dm.View] | None): The data model to load.
+        instance_space (str): The instance space to load the data into.
+        class_by_view_id (dict[ViewId, str] | None): A mapping from view id to class name. Defaults to None.
+        creat_issues (Sequence[NeatIssue] | None): A list of issues that occurred during reading. Defaults to None.
+        tracker (type[Tracker] | None): The tracker to use. Defaults to None.
+    """
+
     def __init__(
         self,
         graph_store: NeatGraphStore,
         data_model: dm.DataModel[dm.View] | None,
         instance_space: str,
         class_by_view_id: dict[ViewId, str] | None = None,
-        creat_issues: Sequence[NeatIssue] | None = None,
+        create_issues: Sequence[NeatIssue] | None = None,
         tracker: type[Tracker] | None = None,
     ):
         super().__init__(graph_store)
         self.data_model = data_model
         self.instance_space = instance_space
         self.class_by_view_id = class_by_view_id or {}
-        self._issues = NeatIssueList[NeatIssue](creat_issues or [])
+        self._issues = NeatIssueList[NeatIssue](create_issues or [])
         self._tracker: type[Tracker] = tracker or LogTracker
 
     @classmethod
@@ -243,11 +255,10 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
         self,
         client: CogniteClient,
         items: list[dm.InstanceApply],
-        return_diffs: bool,
         dry_run: bool,
         read_issues: NeatIssueList,
-    ) -> UploadDiffsID:
-        result = UploadDiffsID(name=type(self).__name__, issues=read_issues)
+    ) -> UploadResult:
+        result = UploadResult[InstanceId](name=type(self).__name__, issues=read_issues)
         try:
             nodes = [item for item in items if isinstance(item, dm.NodeApply)]
             edges = [item for item in items if isinstance(item, dm.EdgeApply)]
@@ -260,16 +271,17 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
             )
         except CogniteAPIError as e:
             result.error_messages.append(str(e))
-            result.failed.append([repr(instance.as_id()) for instance in items])  # type: ignore[arg-type, attr-defined]
+            result.failed_upserted.update(item.as_id() for item in e.failed + e.unknown)
+            result.created.update(item.as_id() for item in e.successful)
         else:
             for instance in itertools.chain(upserted.nodes, upserted.edges):
                 if instance.was_modified and instance.created_time == instance.last_updated_time:
-                    result.created.append(repr(instance.as_id()))
+                    result.created.add(instance.as_id())
                 elif instance.was_modified:
-                    result.changed.append(repr(instance.as_id()))
+                    result.changed.add(instance.as_id())
                 else:
-                    result.unchanged.append(repr(instance.as_id()))
-        return result if return_diffs else result.as_upload_result_ids()  # type: ignore[return-value]
+                    result.unchanged.add(instance.as_id())
+        return result
 
 
 def _triples2dictionary(

{cognite_neat-0.84.0 → cognite_neat-0.85.0}/cognite/neat/graph/stores/_base.py

@@ -14,8 +14,8 @@ from cognite.neat.graph.extractors import RdfFileExtractor, TripleExtractors
 from cognite.neat.graph.models import Triple
 from cognite.neat.graph.queries import Queries
 from cognite.neat.graph.transformers import Transformers
+from cognite.neat.rules.models import InformationRules
 from cognite.neat.rules.models.entities import ClassEntity
-from cognite.neat.rules.models.information import InformationRules
 from cognite.neat.utils.auxiliary import local_import
 
 from ._provenance import Change, Provenance
@@ -66,7 +66,8 @@ class NeatGraphStore:
 
     def add_rules(self, rules: InformationRules) -> None:
         """This method is used to add rules to the graph store and it is the only correct
-        way to add rules to the graph store, after the graph store has been initialized."""
+        way to add rules to the graph store, after the graph store has been initialized.
+        """
 
         self.rules = rules
         self.base_namespace = self.rules.metadata.namespace
@@ -169,7 +170,10 @@ class NeatGraphStore:
         # not yet developed
 
         if not self.rules:
-            warnings.warn("No rules found for the graph store, returning empty list.", stacklevel=2)
+            warnings.warn(
+                "No rules found for the graph store, returning empty list.",
+                stacklevel=2,
+            )
             return []
 
         class_entity = ClassEntity(prefix=self.rules.metadata.prefix, suffix=class_)
@@ -278,3 +282,13 @@ class NeatGraphStore:
                     description=transformer.description,
                 )
             )
+
+    def _repr_html_(self) -> str:
+        provenance = self.provenance._repr_html_()
+
+        return (
+            f"<strong>{type(self).__name__}</strong> A graph store is a container for storing triples. "
+            "It can be queried and transformed to extract information.<br />"
+            "<strong>Provenance</strong> Provenance is a record of changes that have occurred in the graph store.<br />"
+            f"{provenance}"
+        )

{cognite_neat-0.84.0 → cognite_neat-0.85.0}/cognite/neat/graph/stores/_provenance.py

@@ -8,7 +8,6 @@
 
 
 import uuid
-from collections import UserList
 from collections.abc import Sequence
 from dataclasses import dataclass
 from datetime import datetime
@@ -16,6 +15,7 @@ from typing import TypeVar
 
 from rdflib import PROV, RDF, Literal, URIRef
 
+from cognite.neat._shared import FrozenNeatObject, NeatList
 from cognite.neat.constants import DEFAULT_NAMESPACE
 
 
@@ -64,7 +64,7 @@ class Entity:
 
 
 @dataclass(frozen=True)
-class Change:
+class Change(FrozenNeatObject):
     agent: Agent
     activity: Activity
     entity: Entity
@@ -81,11 +81,19 @@ class Change:
         entity = Entity(was_generated_by=activity, was_attributed_to=agent)
         return cls(agent, activity, entity, description)
 
+    def dump(self, aggregate: bool = True) -> dict[str, str]:
+        return {
+            "Agent": self.agent.id_,
+            "Activity": self.activity.id_,
+            "Entity": self.entity.id_,
+            "Description": self.description,
+        }
+
 
 T_Change = TypeVar("T_Change", bound=Change)
 
 
-class Provenance(UserList[T_Change]):
+class Provenance(NeatList[Change]):
     def __init__(self, changes: Sequence[T_Change] | None = None):
         super().__init__(changes or [])