cognite-neat 0.103.1__py3-none-any.whl → 0.105.0__py3-none-any.whl

cognite/neat/_session/_read.py

@@ -1,23 +1,16 @@
-import tempfile
-from datetime import datetime, timezone
-from pathlib import Path
 from typing import Any, Literal, cast
 
 from cognite.client.data_classes.data_modeling import DataModelId, DataModelIdentifier
 
 from cognite.neat._client import NeatClient
-from cognite.neat._constants import COGNITE_SPACES
 from cognite.neat._graph import examples as instances_examples
 from cognite.neat._graph import extractors
 from cognite.neat._issues import IssueList
 from cognite.neat._issues.errors import NeatValueError
+from cognite.neat._issues.warnings import MissingCogniteClientWarning
 from cognite.neat._rules import catalog, importers
-from cognite.neat._rules._shared import ReadRules
 from cognite.neat._rules.importers import BaseImporter
-from cognite.neat._store._provenance import Activity as ProvenanceActivity
-from cognite.neat._store._provenance import Change
-from cognite.neat._store._provenance import Entity as ProvenanceEntity
-from cognite.neat._utils.reader import GitHubReader, HttpFileReader, NeatReader, PathReader
+from cognite.neat._utils.reader import NeatReader
 
 from ._state import SessionState
 from ._wizard import NeatObjectType, RDFFileType, XMLFileType, object_wizard, rdf_dm_wizard, xml_format_wizard
@@ -47,26 +40,6 @@ class BaseReadAPI:
         self._verbose = verbose
         self._client = client
 
-    def _store_rules(self, rules: ReadRules, change: Change) -> IssueList:
-        if self._verbose:
-            if rules.issues.has_errors:
-                print("Data model read failed")
-            else:
-                print("Data model read passed")
-
-        if rules.rules:
-            self._state.data_model.write(rules, change)
-
-        return rules.issues
-
-    def _return_filepath(self, io: Any) -> Path:
-        if isinstance(io, str):
-            return Path(io)
-        elif isinstance(io, Path):
-            return io
-        else:
-            raise NeatValueError(f"Expected str or Path, got {type(io)}")
-
 
 @session_class_wrapper
 class CDFReadAPI(BaseReadAPI):
@@ -103,31 +76,8 @@ class CDFReadAPI(BaseReadAPI):
         if not data_model_id.version:
             raise NeatSessionError("Data model version is required to read a data model.")
 
-        # actual reading of data model
-        start = datetime.now(timezone.utc)
         importer = importers.DMSImporter.from_data_model_id(self._get_client, data_model_id)
-        rules = importer.to_rules()
-        end = datetime.now(timezone.utc)
-
-        # provenance information
-        source_entity = ProvenanceEntity.from_data_model_id(data_model_id)
-        agent = importer.agent
-        activity = ProvenanceActivity(
-            was_associated_with=agent,
-            ended_at_time=end,
-            used=source_entity,
-            started_at_time=start,
-        )
-        target_entity = ProvenanceEntity.from_rules(rules, agent, activity)
-        change = Change(
-            source_entity=source_entity,
-            agent=agent,
-            activity=activity,
-            target_entity=target_entity,
-            description=f"DMS Data model {data_model_id.as_tuple()} read as unverified data model",
-        )
-
-        return self._store_rules(rules, change)
+        return self._state.rule_import(importer)
 
 
 @session_class_wrapper
@@ -215,24 +165,8 @@ class ExcelReadAPI(BaseReadAPI):
             io: file path to the Excel sheet
         """
         reader = NeatReader.create(io)
-        start = datetime.now(timezone.utc)
-        if not isinstance(reader, PathReader):
-            raise NeatValueError("Only file paths are supported for Excel files")
-        importer: importers.ExcelImporter = importers.ExcelImporter(reader.path)
-        input_rules: ReadRules = importer.to_rules()
-        end = datetime.now(timezone.utc)
-
-        if input_rules.rules:
-            change = Change.from_rules_activity(
-                input_rules,
-                importer.agent,
-                start,
-                end,
-                description=f"Excel file {reader!s} read as unverified data model",
-            )
-            self._store_rules(input_rules, change)
-        self._state.data_model.issue_lists.append(input_rules.issues)
-        return input_rules.issues
+        path = reader.materialize_path()
+        return self._state.rule_import(importers.ExcelImporter(path))
 
 
 @session_class_wrapper
@@ -241,85 +175,43 @@ class ExcelExampleAPI(BaseReadAPI):
 
     @property
     def pump_example(self) -> IssueList:
-        """Reads the Nordic 44 knowledge graph into the NeatSession graph store."""
-        start = datetime.now(timezone.utc)
+        """Reads the Hello World pump example into the NeatSession."""
         importer: importers.ExcelImporter = importers.ExcelImporter(catalog.hello_world_pump)
-        input_rules: ReadRules = importer.to_rules()
-        end = datetime.now(timezone.utc)
-
-        if input_rules.rules:
-            change = Change.from_rules_activity(
-                input_rules,
-                importer.agent,
-                start,
-                end,
-                description="Pump Example read as unverified data model",
-            )
-            self._store_rules(input_rules, change)
-        self._state.data_model.issue_lists.append(input_rules.issues)
-        return input_rules.issues
+        return self._state.rule_import(importer)
 
 
 @session_class_wrapper
 class YamlReadAPI(BaseReadAPI):
-    """Reads a yaml with either neat rules, or several toolkit yaml files to import Data Model(s) into NeatSession.
-
-    Args:
-        io: file path to the Yaml file in the case of "neat" yaml, or path to a zip folder or directory with several
-        Yaml files in the case of "toolkit".
+    def __call__(self, io: Any, format: Literal["neat", "toolkit"] = "neat") -> IssueList:
+        """Reads a yaml with either neat rules, or several toolkit yaml files to import Data Model(s) into NeatSession.
 
-    Example:
-        ```python
-        neat.read.yaml("path_to_toolkit_yamls")
-        ```
-    """
+        Args:
+            io: File path to the Yaml file in the case of "neat" yaml, or path to a zip folder or directory with several
+                Yaml files in the case of "toolkit".
+            format: The format of the yaml file(s). Can be either "neat" or "toolkit".
 
-    def __call__(self, io: Any, format: Literal["neat", "toolkit"] = "neat") -> IssueList:
+        Example:
+            ```python
+            neat.read.yaml("path_to_toolkit_yamls")
+            ```
+        """
         reader = NeatReader.create(io)
-        if not isinstance(reader, PathReader):
-            raise NeatValueError("Only file paths are supported for YAML files")
-        start = datetime.now(timezone.utc)
+        path = reader.materialize_path()
         importer: BaseImporter
         if format == "neat":
-            importer = importers.YAMLImporter.from_file(reader.path)
+            importer = importers.YAMLImporter.from_file(path, source_name=f"{reader!s}")
         elif format == "toolkit":
-            if reader.path.is_file():
-                dms_importer = importers.DMSImporter.from_zip_file(reader.path)
-            elif reader.path.is_dir():
-                dms_importer = importers.DMSImporter.from_directory(reader.path)
-            else:
-                raise NeatValueError(f"Unsupported YAML format: {format}")
-            ref_containers = dms_importer.root_schema.referenced_container()
-            if system_container_ids := [
-                container_id for container_id in ref_containers if container_id.space in COGNITE_SPACES
-            ]:
-                if self._client is None:
-                    raise NeatSessionError(
-                        "No client provided. You are referencing Cognite containers in your data model, "
-                        "NEAT needs a client to lookup the container definitions. "
-                        "Please set the client in the session, NeatSession(client=client)."
-                    )
-                system_containers = self._client.loaders.containers.retrieve(system_container_ids)
-                dms_importer.update_referenced_containers(system_containers)
-
+            dms_importer = importers.DMSImporter.from_path(path, self._client)
+            if dms_importer.issue_list.has_warning_type(MissingCogniteClientWarning):
+                raise NeatSessionError(
+                    "No client provided. You are referencing Cognite containers in your data model, "
+                    "NEAT needs a client to lookup the container definitions. "
+                    "Please set the client in the session, NeatSession(client=client)."
+                )
             importer = dms_importer
         else:
             raise NeatValueError(f"Unsupported YAML format: {format}")
-        input_rules: ReadRules = importer.to_rules()
-
-        end = datetime.now(timezone.utc)
-
-        if input_rules.rules:
-            change = Change.from_rules_activity(
-                input_rules,
-                importer.agent,
-                start,
-                end,
-                description=f"YAML file {reader!s} read as unverified data model",
-            )
-            self._store_rules(input_rules, change)
-
-        return input_rules.issues
+        return self._state.rule_import(importer)
 
 
 @session_class_wrapper
@@ -342,17 +234,9 @@ class CSVReadAPI(BaseReadAPI):
     """
 
     def __call__(self, io: Any, type: str, primary_key: str) -> None:
-        reader = NeatReader.create(io)
-        if isinstance(reader, HttpFileReader):
-            path = Path(tempfile.gettempdir()).resolve() / reader.name
-            path.write_text(reader.read_text(), encoding="utf-8", newline="\n")
-        elif isinstance(reader, PathReader):
-            path = reader.path
-        else:
-            raise NeatValueError("Only file paths are supported for CSV files")
         engine = import_engine()
         engine.set.format = "csv"
-        engine.set.file = path
+        engine.set.file = NeatReader.create(io).materialize_path()
         engine.set.type = type
         engine.set.primary_key = primary_key
         extractor = engine.create_extractor()
@@ -374,14 +258,7 @@ class XMLReadAPI(BaseReadAPI):
         io: Any,
         format: XMLFileType | None = None,
     ) -> None:
-        reader = NeatReader.create(io)
-        if isinstance(reader, GitHubReader):
-            path = Path(tempfile.gettempdir()).resolve() / reader.name
-            path.write_text(reader.read_text())
-        elif isinstance(reader, PathReader):
-            path = reader.path
-        else:
-            raise NeatValueError("Only file paths are supported for XML files")
+        path = NeatReader.create(io).materialize_path()
         if format is None:
             format = xml_format_wizard()
 
@@ -394,7 +271,7 @@ class XMLReadAPI(BaseReadAPI):
         else:
             raise NeatValueError("Only support XML files of DEXPI format at the moment.")
 
-    def dexpi(self, path):
+    def dexpi(self, io: Any) -> None:
         """Reads a DEXPI file into the NeatSession.
 
         Args:
@@ -405,13 +282,14 @@ class XMLReadAPI(BaseReadAPI):
             neat.read.xml.dexpi("url_or_path_to_dexpi_file")
             ```
         """
+        path = NeatReader.create(io).materialize_path()
         engine = import_engine()
         engine.set.format = "dexpi"
         engine.set.file = path
         extractor = engine.create_extractor()
         self._state.instances.store.write(extractor)
 
-    def aml(self, path):
+    def aml(self, io: Any):
         """Reads an AML file into NeatSession.
 
         Args:
@@ -422,6 +300,7 @@ class XMLReadAPI(BaseReadAPI):
             neat.read.xml.aml("url_or_path_to_aml_file")
             ```
         """
+        path = NeatReader.create(io).materialize_path()
         engine = import_engine()
         engine.set.format = "aml"
         engine.set.file = path
@@ -452,25 +331,9 @@ class RDFReadAPI(BaseReadAPI):
             neat.read.rdf.ontology("url_or_path_to_owl_source")
             ```
         """
-        start = datetime.now(timezone.utc)
         reader = NeatReader.create(io)
-        if not isinstance(reader, PathReader):
-            raise NeatValueError("Only file paths are supported for RDF files")
-        importer = importers.OWLImporter.from_file(reader.path)
-        input_rules: ReadRules = importer.to_rules()
-        end = datetime.now(timezone.utc)
-
-        if input_rules.rules:
-            change = Change.from_rules_activity(
-                input_rules,
-                importer.agent,
-                start,
-                end,
-                description=f"Ontology file {reader!s} read as unverified data model",
-            )
-            self._store_rules(input_rules, change)
-
-        return input_rules.issues
+        importer = importers.OWLImporter.from_file(reader.materialize_path(), source_name=f"file {reader!s}")
+        return self._state.rule_import(importer)
 
     def imf(self, io: Any) -> IssueList:
         """Reads IMF Types provided as SHACL shapes into NeatSession.
@@ -483,25 +346,9 @@ class RDFReadAPI(BaseReadAPI):
             neat.read.rdf.imf("url_or_path_to_imf_source")
             ```
         """
-        start = datetime.now(timezone.utc)
         reader = NeatReader.create(io)
-        if not isinstance(reader, PathReader):
-            raise NeatValueError("Only file paths are supported for RDF files")
-        importer = importers.IMFImporter.from_file(reader.path)
-        input_rules: ReadRules = importer.to_rules()
-        end = datetime.now(timezone.utc)
-
-        if input_rules.rules:
-            change = Change.from_rules_activity(
-                input_rules,
-                importer.agent,
-                start,
-                end,
-                description=f"IMF Types file {reader!s} read as unverified data model",
-            )
-            self._store_rules(input_rules, change)
-
-        return input_rules.issues
+        importer = importers.IMFImporter.from_file(reader.materialize_path(), source_name=f"file {reader!s}")
+        return self._state.rule_import(importer)
 
     def __call__(
         self,
@@ -527,10 +374,7 @@ class RDFReadAPI(BaseReadAPI):
 
         elif type == "instances":
             reader = NeatReader.create(io)
-            if not isinstance(reader, PathReader):
-                raise NeatValueError("Only file paths are supported for RDF files")
-
-            self._state.instances.store.write(extractors.RdfFileExtractor(reader.path))
+            self._state.instances.store.write(extractors.RdfFileExtractor(reader.materialize_path()))
             return IssueList()
         else:
             raise NeatSessionError(f"Expected data model or instances, got {type}")

cognite/neat/_session/_set.py

@@ -1,10 +1,9 @@
-from datetime import datetime, timezone
-
 from cognite.client import data_modeling as dm
 
 from cognite.neat._constants import COGNITE_MODELS
+from cognite.neat._issues import IssueList
+from cognite.neat._rules.models import DMSRules
 from cognite.neat._rules.transformers import SetIDDMSModel
-from cognite.neat._store._provenance import Change
 
 from ._state import SessionState
 from .exceptions import NeatSessionError, session_class_wrapper
@@ -18,7 +17,7 @@ class SetAPI:
         self._state = state
         self._verbose = verbose
 
-    def data_model_id(self, new_model_id: dm.DataModelId | tuple[str, str, str]) -> None:
+    def data_model_id(self, new_model_id: dm.DataModelId | tuple[str, str, str]) -> IssueList:
         """Sets the data model ID of the latest verified data model. Set the data model id as a tuple of strings
         following the template (<data_model_space>, <data_model_name>, <data_model_version>).
 
@@ -28,34 +27,11 @@ class SetAPI:
             neat.set.data_model_id(("my_data_model_space", "My_Data_Model", "v1"))
             ```
         """
-        if res := self._state.data_model.last_verified_dms_rules:
-            source_id, rules = res
-
+        rules = self._state.rule_store.get_last_successful_entity().result
+        if isinstance(rules, DMSRules):
             if rules.metadata.as_data_model_id() in COGNITE_MODELS:
                 raise NeatSessionError(
                     "Cannot change the data model ID of a Cognite Data Model in NeatSession"
                     " due to temporarily issue with the reverse direct relation interpretation"
                 )
-
-            start = datetime.now(timezone.utc)
-            transformer = SetIDDMSModel(new_model_id)
-
-            output = transformer.transform(rules)
-            end = datetime.now(timezone.utc)
-
-            # Provenance
-            change = Change.from_rules_activity(
-                output.rules,
-                transformer.agent,
-                start,
-                end,
-                "Changed data model id",
-                self._state.data_model.provenance.source_entity(source_id)
-                or self._state.data_model.provenance.target_entity(source_id),
-            )
-
-            self._state.data_model.write(output.rules, change)
-            if self._verbose:
-                print(f"Data model ID set to {new_model_id}")
-        else:
-            print("No verified DMS data model available")
+        return self._state.rule_transform(SetIDDMSModel(new_model_id))

cognite/neat/_session/_show.py

@@ -5,7 +5,6 @@ from typing import Any, cast
 import networkx as nx
 from IPython.display import HTML, display
 from pyvis.network import Network as PyVisNetwork  # type: ignore
-from rdflib import URIRef
 
 from cognite.neat._constants import IN_NOTEBOOK, IN_PYODIDE
 from cognite.neat._rules._constants import EntityTypes
@@ -14,7 +13,7 @@ from cognite.neat._rules.models.entities._single_value import ClassEntity, ViewE
 from cognite.neat._rules.models.information._rules import InformationRules
 from cognite.neat._session.exceptions import NeatSessionError
 from cognite.neat._utils.io_ import to_directory_compatible
-from cognite.neat._utils.rdf_ import remove_namespace_from_uri
+from cognite.neat._utils.rdf_ import remove_namespace_from_uri, uri_display_name
 
 from ._state import SessionState
 from .exceptions import session_class_wrapper
@@ -99,12 +98,12 @@ class ShowDataModelAPI(ShowBaseAPI):
         self.implements = ShowDataModelImplementsAPI(self._state)
 
     def __call__(self) -> Any:
-        if not self._state.data_model.has_verified_rules:
+        if not self._state.rule_store.has_verified_rules:
             raise NeatSessionError(
                 "No verified data model available. Try using [bold].verify()[/bold] to verify data model."
             )
 
-        rules = self._state.data_model.last_verified_rule[1]
+        rules = self._state.rule_store.last_verified_rule
 
         if isinstance(rules, DMSRules):
             di_graph = self._generate_dms_di_graph(rules)
@@ -188,12 +187,12 @@ class ShowDataModelImplementsAPI(ShowBaseAPI):
         self._state = state
 
     def __call__(self) -> Any:
-        if not self._state.data_model.has_verified_rules:
+        if not self._state.rule_store.has_verified_rules:
             raise NeatSessionError(
                 "No verified data model available. Try using [bold].verify()[/bold] to verify data model."
             )
 
-        rules = self._state.data_model.last_verified_rule[1]
+        rules = self._state.rule_store.last_verified_rule
 
         if isinstance(rules, DMSRules):
             di_graph = self._generate_dms_di_graph(rules)
@@ -270,20 +269,21 @@ class ShowDataModelProvenanceAPI(ShowBaseAPI):
         self._state = state
 
     def __call__(self) -> Any:
-        if not self._state.data_model.provenance:
+        if not self._state.rule_store.provenance:
             raise NeatSessionError("No data model available. Try using [bold].read[/bold] to load data model.")
 
         di_graph = self._generate_dm_provenance_di_graph_and_types()
-        return self._generate_visualization(di_graph, name="data_model_provenance.html")
+        unique_hash = self._state.rule_store.calculate_provenance_hash(shorten=True)
+        return self._generate_visualization(di_graph, name=f"data_model_provenance_{unique_hash}.html")
 
     def _generate_dm_provenance_di_graph_and_types(self) -> nx.DiGraph:
         di_graph = nx.DiGraph()
-        hex_colored_types = _generate_hex_color_per_type(["Agent", "Entity", "Activity"])
+        hex_colored_types = _generate_hex_color_per_type(["Agent", "Entity", "Activity", "Export", "Pruned"])
 
-        for change in self._state.data_model.provenance:
-            source = self._shorten_id(change.source_entity.id_)
-            target = self._shorten_id(change.target_entity.id_)
-            agent = self._shorten_id(change.agent.id_)
+        for change in self._state.rule_store.provenance:
+            source = uri_display_name(change.source_entity.id_)
+            target = uri_display_name(change.target_entity.id_)
+            agent = uri_display_name(change.agent.id_)
 
             di_graph.add_node(
                 source,
@@ -312,15 +312,34 @@ class ShowDataModelProvenanceAPI(ShowBaseAPI):
             di_graph.add_edge(source, agent, label="used", color="grey")
             di_graph.add_edge(agent, target, label="generated", color="grey")
 
-        return di_graph
+        for source_id, exports in self._state.rule_store.exports_by_source_entity_id.items():
+            source_shorten = uri_display_name(source_id)
+            for export in exports:
+                export_id = uri_display_name(export.target_entity.id_)
+                di_graph.add_node(
+                    export_id,
+                    label=export_id,
+                    type="Export",
+                    title="Export",
+                    color=hex_colored_types["Export"],
+                )
+                di_graph.add_edge(source_shorten, export_id, label="exported", color="grey")
 
-    @staticmethod
-    def _shorten_id(thing: URIRef) -> str:
-        if "https://cognitedata.com/dms/data-model/" in thing:
-            return "DMS(" + ",".join(thing.replace("https://cognitedata.com/dms/data-model/", "").split("/")) + ")"
-        elif "http://purl.org/cognite/neat/data-model/" in thing:
-            return "NEAT(" + ",".join(thing.replace("http://purl.org/cognite/neat/data-model/", "").split("/")) + ")"
-        return remove_namespace_from_uri(thing)
+        for pruned_lists in self._state.rule_store.pruned_by_source_entity_id.values():
+            for prune_path in pruned_lists:
+                for change in prune_path:
+                    source = uri_display_name(change.source_entity.id_)
+                    target = uri_display_name(change.target_entity.id_)
+                    di_graph.add_node(
+                        target,
+                        label=target,
+                        type="Pruned",
+                        title="Pruned",
+                        color=hex_colored_types["Pruned"],
+                    )
+                    di_graph.add_edge(source, target, label="pruned", color="grey")
+
+        return di_graph
 
 
 @session_class_wrapper