cognite-neat 0.105.2__py3-none-any.whl → 0.107.0__py3-none-any.whl

cognite/neat/_session/_prepare.py

@@ -4,11 +4,9 @@ from typing import Any, Literal, cast
 from cognite.client.data_classes.data_modeling import DataModelIdentifier
 from rdflib import URIRef
 
-from cognite.neat._client import NeatClient
 from cognite.neat._constants import (
     get_default_prefixes_and_namespaces,
 )
-from cognite.neat._graph import extractors
 from cognite.neat._graph.transformers import (
     AttachPropertyFromTargetToSource,
     ConnectionToLiteral,
@@ -25,7 +23,6 @@ from cognite.neat._issues import IssueList
 from cognite.neat._issues.errors import NeatValueError
 from cognite.neat._rules.models.dms import DMSValidation
 from cognite.neat._rules.transformers import (
-    AddClassImplements,
     IncludeReferenced,
     PrefixEntities,
     ReduceCogniteModel,
@@ -47,10 +44,10 @@ class PrepareAPI:
     inferring a data model or exporting the knowledge graph to a desired destination.
     """
 
-    def __init__(self, client: NeatClient | None, state: SessionState, verbose: bool) -> None:
+    def __init__(self, state: SessionState, verbose: bool) -> None:
         self._state = state
         self._verbose = verbose
-        self.data_model = DataModelPrepareAPI(client, state, verbose)
+        self.data_model = DataModelPrepareAPI(state, verbose)
         self.instances = InstancePrepareAPI(state, verbose)
 
 
@@ -349,41 +346,6 @@ class InstancePrepareAPI:
         transformer = ConnectionToLiteral(subject_type, subject_predicate)
         self._state.instances.store.transform(transformer)
 
-    def classic_to_core(self) -> None:
-        """Prepares extracted CDF classic graph for the Core Data model.
-
-        !!! note "This method bundles several graph transformers which"
-            - Convert relationships to edges
-            - Convert TimeSeries.type from bool to enum
-            - Convert all properties 'source' to a connection to SourceSystem
-            - Convert all properties 'labels' from a connection to a string
-
-        Example:
-            Apply classic to core transformations:
-            ```python
-            neat.prepare.instances.classic_to_core()
-            ```
-        """
-        self.relationships_as_edges()
-        self.convert_data_type(
-            ("TimeSeries", "isString"), convert=lambda is_string: "string" if is_string else "numeric"
-        )
-        self.property_to_type((None, "source"), "SourceSystem", "name")
-        for type_ in [
-            extractors.EventsExtractor._default_rdf_type,
-            extractors.AssetsExtractor._default_rdf_type,
-            extractors.FilesExtractor._default_rdf_type,
-        ]:
-            try:
-                subject_type, subject_predicate = self._get_type_and_property_uris(type_, "labels")
-            except NeatValueError:
-                # If the type_.labels does not exist, continue. This is not an error, it just means that the
-                # Labels is not used in the graph for that type.
-                continue
-            else:
-                transformer = ConnectionToLiteral(subject_type, subject_predicate)
-                self._state.instances.store.transform(transformer)
-
 
 @session_class_wrapper
 class DataModelPrepareAPI:
@@ -391,8 +353,7 @@ class DataModelPrepareAPI:
     to a desired destination.
     """
 
-    def __init__(self, client: NeatClient | None, state: SessionState, verbose: bool) -> None:
-        self._client = client
+    def __init__(self, state: SessionState, verbose: bool) -> None:
         self._state = state
         self._verbose = verbose
 
@@ -507,14 +468,15 @@ class DataModelPrepareAPI:
             self._state.rule_store.last_verified_dms_rules
         ).imported_views_and_containers_ids()
         transformers: list[RulesTransformer] = []
-        if (view_ids or container_ids) and self._client is None:
+        client = self._state.client
+        if (view_ids or container_ids) and client is None:
             raise NeatSessionError(
                 "No client provided. You are referencing unknown views and containers in your data model, "
                 "NEAT needs a client to lookup the definitions. "
                 "Please set the client in the session, NeatSession(client=client)."
             )
-        elif (view_ids or container_ids) and self._client:
-            transformers.append(IncludeReferenced(self._client, include_properties=True))
+        elif (view_ids or container_ids) and client:
+            transformers.append(IncludeReferenced(client, include_properties=True))
 
         transformers.append(
             ToDataProductModel(
@@ -539,20 +501,10 @@ class DataModelPrepareAPI:
 
     def include_referenced(self) -> IssueList:
         """Include referenced views and containers in the data model."""
-        if self._client is None:
+        if self._state.client is None:
             raise NeatSessionError(
                 "No client provided. You are referencing unknown views and containers in your data model, "
                 "NEAT needs a client to lookup the definitions. "
                 "Please set the client in the session, NeatSession(client=client)."
             )
-        return self._state.rule_transform(IncludeReferenced(self._client))
-
-    def add_implements_to_classes(self, suffix: Literal["Edge"], implements: str = "Edge") -> IssueList:
-        """All classes with the suffix will have the implements property set to the given value.
-
-        Args:
-            suffix: The suffix of the classes to add the implements property to.
-            implements:  The value of the implements property to set.
-
-        """
-        return self._state.rule_transform(AddClassImplements(implements, suffix))
+        return self._state.rule_transform(IncludeReferenced(self._state.client))

cognite/neat/_session/_read.py

@@ -1,15 +1,19 @@
 from typing import Any, Literal, cast
 
 from cognite.client.data_classes.data_modeling import DataModelId, DataModelIdentifier
+from cognite.client.utils.useful_types import SequenceNotStr
 
 from cognite.neat._client import NeatClient
+from cognite.neat._constants import CLASSIC_CDF_NAMESPACE
 from cognite.neat._graph import examples as instances_examples
 from cognite.neat._graph import extractors
+from cognite.neat._graph.transformers import ConvertLiteral, LiteralToEntity, LookupRelationshipSourceTarget
 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.importers import BaseImporter
+from cognite.neat._rules.transformers import ClassicPrepareCore
 from cognite.neat._utils.reader import NeatReader
 
 from ._state import SessionState
@@ -22,23 +26,38 @@ from .exceptions import NeatSessionError, session_class_wrapper
 class ReadAPI:
     """Read from a data source into NeatSession graph store."""
 
-    def __init__(self, state: SessionState, client: NeatClient | None, verbose: bool) -> None:
+    def __init__(self, state: SessionState, verbose: bool) -> None:
         self._state = state
         self._verbose = verbose
-        self.cdf = CDFReadAPI(state, client, verbose)
-        self.rdf = RDFReadAPI(state, client, verbose)
-        self.excel = ExcelReadAPI(state, client, verbose)
-        self.csv = CSVReadAPI(state, client, verbose)
-        self.yaml = YamlReadAPI(state, client, verbose)
-        self.xml = XMLReadAPI(state, client, verbose)
+        self.cdf = CDFReadAPI(state, verbose)
+        self.rdf = RDFReadAPI(state, verbose)
+        self.excel = ExcelReadAPI(state, verbose)
+        self.csv = CSVReadAPI(state, verbose)
+        self.yaml = YamlReadAPI(state, verbose)
+        self.xml = XMLReadAPI(state, verbose)
+
+    def session(self, io: Any) -> None:
+        """Reads a Neat Session from a zip file.
+
+        Args:
+            io: file path to the Neat Session
+
+        Example:
+            ```python
+            neat.read.session("path_to_neat_session")
+            ```
+        """
+        reader = NeatReader.create(io)
+        path = reader.materialize_path()
+
+        self._state.instances.store.write(extractors.RdfFileExtractor.from_zip(path))
 
 
 @session_class_wrapper
 class BaseReadAPI:
-    def __init__(self, state: SessionState, client: NeatClient | None, verbose: bool) -> None:
+    def __init__(self, state: SessionState, verbose: bool) -> None:
         self._state = state
         self._verbose = verbose
-        self._client = client
 
 
 @session_class_wrapper
@@ -48,15 +67,15 @@ class CDFReadAPI(BaseReadAPI):
 
     """
 
-    def __init__(self, state: SessionState, client: NeatClient | None, verbose: bool) -> None:
-        super().__init__(state, client, verbose)
-        self.classic = CDFClassicAPI(state, client, verbose)
+    def __init__(self, state: SessionState, verbose: bool) -> None:
+        super().__init__(state, verbose)
+        self.classic = CDFClassicAPI(state, verbose)
 
     @property
     def _get_client(self) -> NeatClient:
-        if self._client is None:
+        if self._state.client is None:
             raise NeatValueError("No client provided. Please provide a client to read a data model.")
-        return self._client
+        return self._state.client
 
     def data_model(self, data_model_id: DataModelIdentifier) -> IssueList:
         """Reads a Data Model from CDF to the knowledge graph.
@@ -79,6 +98,24 @@ class CDFReadAPI(BaseReadAPI):
         importer = importers.DMSImporter.from_data_model_id(self._get_client, data_model_id)
         return self._state.rule_import(importer)
 
+    def graph(
+        self, data_model_id: DataModelIdentifier, instance_space: str | SequenceNotStr[str] | None = None
+    ) -> IssueList:
+        """Reads a knowledge graph from Cognite Data Fusion (CDF).
+
+        Args:
+            data_model_id: Tuple of strings with the id of a CDF Data Model.
+            instance_space: The instance spaces to extract. If None, all instance spaces are extracted.
+
+        Returns:
+            IssueList: A list of issues that occurred during the extraction.
+
+        """
+        extractor = extractors.DMSGraphExtractor.from_data_model_id(
+            data_model_id, self._get_client, instance_space=instance_space
+        )
+        return self._state.write_graph(extractor)
+
 
 @session_class_wrapper
 class CDFClassicAPI(BaseReadAPI):
@@ -89,11 +126,11 @@ class CDFClassicAPI(BaseReadAPI):
 
     @property
     def _get_client(self) -> NeatClient:
-        if self._client is None:
+        if self._state.client is None:
             raise ValueError("No client provided. Please provide a client to read a data model.")
-        return self._client
+        return self._state.client
 
-    def graph(self, root_asset_external_id: str, limit_per_type: int | None = None) -> None:
+    def graph(self, root_asset_external_id: str, limit_per_type: int | None = None) -> IssueList:
         """Reads the classic knowledge graph from CDF.
 
         The Classic Graph consists of the following core resource type.
@@ -129,14 +166,45 @@ class CDFClassicAPI(BaseReadAPI):
             root_asset_external_id: The external id of the root asset
             limit_per_type: The maximum number of nodes to extract per core node type. If None, all nodes are extracted.
 
+        Returns:
+            IssueList: A list of issues that occurred during the extraction.
+
+        Example:
+            ```python
+            neat.read.cdf.graph("root_asset_external_id")
+            ```
         """
+        namespace = CLASSIC_CDF_NAMESPACE
         extractor = extractors.ClassicGraphExtractor(
-            self._get_client, root_asset_external_id=root_asset_external_id, limit_per_type=limit_per_type
+            self._get_client,
+            root_asset_external_id=root_asset_external_id,
+            limit_per_type=limit_per_type,
+            namespace=namespace,
+            prefix="Classic",
         )
-
-        self._state.instances.store.write(extractor)
-        if self._verbose:
-            print(f"Classic Graph {root_asset_external_id} read successfully")
+        issues = self._state.write_graph(extractor)
+        issues.action = "Read Classic Graph"
+        if issues:
+            print("Use the .inspect.issues() for more details")
+
+        # Converting the instances from classic to core
+        self._state.instances.store.transform(LookupRelationshipSourceTarget(namespace, "Classic"))
+        self._state.instances.store.transform(
+            ConvertLiteral(
+                namespace["ClassicTimeSeries"],
+                namespace["isString"],
+                lambda is_string: "string" if is_string else "numeric",
+            )
+        )
+        self._state.instances.store.transform(
+            LiteralToEntity(None, namespace["source"], "ClassicSourceSystem", "name"),
+        )
+        # Updating the information model.
+        self._state.rule_store.transform(ClassicPrepareCore(namespace))
+        # Update the instance store with the latest rules
+        information_rules = self._state.rule_store.last_verified_information_rules
+        self._state.instances.store.rules = information_rules
+        return issues
 
 
 @session_class_wrapper
@@ -153,9 +221,9 @@ class ExcelReadAPI(BaseReadAPI):
         ```
     """
 
-    def __init__(self, state: SessionState, client: NeatClient | None, verbose: bool) -> None:
-        super().__init__(state, client, verbose)
-        self.examples = ExcelExampleAPI(state, client, verbose)
+    def __init__(self, state: SessionState, verbose: bool) -> None:
+        super().__init__(state, verbose)
+        self.examples = ExcelExampleAPI(state, verbose)
 
     def __call__(self, io: Any) -> IssueList:
         """Reads a Neat Excel Rules sheet to the graph store. The rules sheet may stem from an Information architect,
@@ -201,7 +269,7 @@ class YamlReadAPI(BaseReadAPI):
         if format == "neat":
             importer = importers.YAMLImporter.from_file(path, source_name=f"{reader!s}")
         elif format == "toolkit":
-            dms_importer = importers.DMSImporter.from_path(path, self._client)
+            dms_importer = importers.DMSImporter.from_path(path, self._state.client)
             if dms_importer.issue_list.has_warning_type(MissingCogniteClientWarning):
                 raise NeatSessionError(
                     "No client provided. You are referencing Cognite containers in your data model, "
@@ -316,8 +384,8 @@ class RDFReadAPI(BaseReadAPI):
         io: file path or url to the RDF source
     """
 
-    def __init__(self, state: SessionState, client: NeatClient | None, verbose: bool) -> None:
-        super().__init__(state, client, verbose)
+    def __init__(self, state: SessionState, verbose: bool) -> None:
+        super().__init__(state, verbose)
         self.examples = RDFExamples(state)
 
     def ontology(self, io: Any) -> IssueList:
@@ -387,7 +455,6 @@ class RDFExamples:
     def __init__(self, state: SessionState) -> None:
         self._state = state
 
-    @property
     def nordic44(self) -> IssueList:
         """Reads the Nordic 44 knowledge graph into the NeatSession graph store."""
         self._state.instances.store.write(extractors.RdfFileExtractor(instances_examples.nordic44_knowledge_graph))

cognite/neat/_session/_set.py

@@ -1,5 +1,7 @@
+from cognite.client import CogniteClient
 from cognite.client import data_modeling as dm
 
+from cognite.neat._client import NeatClient
 from cognite.neat._constants import COGNITE_MODELS
 from cognite.neat._issues import IssueList
 from cognite.neat._rules.models import DMSRules
@@ -35,3 +37,10 @@ class SetAPI:
                     " due to temporarily issue with the reverse direct relation interpretation"
                 )
         return self._state.rule_transform(SetIDDMSModel(new_model_id))
+
+    def client(self, client: CogniteClient) -> None:
+        """Sets the client to be used in the session."""
+        self._state.client = NeatClient(client)
+        if self._verbose:
+            print(f"Client set to {self._state.client.config.project} CDF project.")
+        return None

cognite/neat/_session/_state.py

@@ -1,6 +1,8 @@
 from dataclasses import dataclass, field
 from typing import Literal, cast
 
+from cognite.neat._client import NeatClient
+from cognite.neat._graph.extractors import KnowledgeGraphExtractor
 from cognite.neat._issues import IssueList
 from cognite.neat._rules.importers import BaseImporter, InferenceImporter
 from cognite.neat._rules.models import DMSRules, InformationRules
@@ -15,10 +17,11 @@ from .exceptions import NeatSessionError
 
 
 class SessionState:
-    def __init__(self, store_type: Literal["memory", "oxigraph"]) -> None:
+    def __init__(self, store_type: Literal["memory", "oxigraph"], client: NeatClient | None = None) -> None:
         self.instances = InstancesState(store_type)
         self.rule_store = NeatRulesStore()
         self.last_reference: DMSRules | InformationRules | None = None
+        self.client = client
 
     def rule_transform(self, *transformer: RulesTransformer) -> IssueList:
         if not transformer:
@@ -60,6 +63,12 @@ class SessionState:
             issues.hint = "Use the .inspect.issues() for more details."
         return issues
 
+    def write_graph(self, extractor: KnowledgeGraphExtractor) -> IssueList:
+        self.instances.store.write(extractor)
+        issue_list = self.rule_store.import_graph(extractor)
+        self.instances.store.add_rules(self.rule_store.last_verified_information_rules)
+        return issue_list
+
 
 @dataclass
 class InstancesState:

cognite/neat/_session/_to.py

@@ -1,10 +1,11 @@
+import warnings
+import zipfile
 from collections.abc import Collection
 from pathlib import Path
 from typing import Any, Literal, overload
 
 from cognite.client.data_classes.data_modeling import SpaceApply
 
-from cognite.neat._client import NeatClient
 from cognite.neat._constants import COGNITE_MODELS
 from cognite.neat._graph import loaders
 from cognite.neat._rules import exporters
@@ -25,10 +26,10 @@ class ToAPI:
 
     """
 
-    def __init__(self, state: SessionState, client: NeatClient | None, verbose: bool) -> None:
+    def __init__(self, state: SessionState, verbose: bool) -> None:
         self._state = state
         self._verbose = verbose
-        self.cdf = CDFToAPI(state, client, verbose)
+        self.cdf = CDFToAPI(state, verbose)
 
     def excel(
         self,
@@ -81,11 +82,47 @@ class ToAPI:
         exporter = exporters.ExcelExporter(styling="maximal", reference_rules_with_prefix=reference_rules_with_prefix)
         return self._state.rule_store.export_to_file(exporter, Path(io))
 
+    def session(self, io: Any) -> None:
+        """Export the current session to a file.
+
+        Args:
+            io: The file path to file-like object to write the session to.
+
+        Example:
+            Export the session to a file
+            ```python
+            session_file_name = "neat_session.zip"
+            neat.to.session(session_file_name)
+            ```
+        """
+        filepath = Path(io)
+        if filepath.suffix not in {".zip"}:
+            warnings.warn("File extension is not .zip, adding it to the file name", stacklevel=2)
+            filepath = filepath.with_suffix(".zip")
+
+        filepath.parent.mkdir(exist_ok=True, parents=True)
+
+        with zipfile.ZipFile(filepath, "w") as zip_ref:
+            zip_ref.writestr(
+                "neat-session/instances/instances.ttl",
+                self._state.instances.store.serialize(),
+            )
+
     @overload
-    def yaml(self, io: None, format: Literal["neat"] = "neat", skip_system_spaces: bool = True) -> str: ...
+    def yaml(
+        self,
+        io: None,
+        format: Literal["neat"] = "neat",
+        skip_system_spaces: bool = True,
+    ) -> str: ...
 
     @overload
-    def yaml(self, io: Any, format: Literal["neat", "toolkit"] = "neat", skip_system_spaces: bool = True) -> None: ...
+    def yaml(
+        self,
+        io: Any,
+        format: Literal["neat", "toolkit"] = "neat",
+        skip_system_spaces: bool = True,
+    ) -> None: ...
 
     def yaml(
         self, io: Any | None = None, format: Literal["neat", "toolkit"] = "neat", skip_system_spaces: bool = True
@@ -153,8 +190,7 @@ class ToAPI:
 class CDFToAPI:
     """Write a verified Data Model and Instances to CDF."""
 
-    def __init__(self, state: SessionState, client: NeatClient | None, verbose: bool) -> None:
-        self._client = client
+    def __init__(self, state: SessionState, verbose: bool) -> None:
         self._state = state
         self._verbose = verbose
 
@@ -166,9 +202,9 @@ class CDFToAPI:
             Note this space is required to be different than the space with the data model.
 
         """
-        if not self._client:
+        if not self._state.client:
             raise NeatSessionError("No CDF client provided!")
-
+        client = self._state.client
         space = space or f"{self._state.rule_store.last_verified_dms_rules.metadata.space}_instances"
 
         if space and space == self._state.rule_store.last_verified_dms_rules.metadata.space:
@@ -176,16 +212,16 @@ class CDFToAPI:
         elif not PATTERNS.space_compliance.match(str(space)):
             raise NeatSessionError("Please provide a valid space name. {PATTERNS.space_compliance.pattern}")
 
-        if not self._client.data_modeling.spaces.retrieve(space):
-            self._client.data_modeling.spaces.apply(SpaceApply(space=space))
+        if not client.data_modeling.spaces.retrieve(space):
+            client.data_modeling.spaces.apply(SpaceApply(space=space))
 
         loader = loaders.DMSLoader.from_rules(
             self._state.rule_store.last_verified_dms_rules,
             self._state.instances.store,
             instance_space=space,
-            client=self._client,
+            client=client,
         )
-        result = loader.load_into_cdf(self._client)
+        result = loader.load_into_cdf(client)
         self._state.instances.outcome.append(result)
         print("You can inspect the details with the .inspect.outcome.instances(...) method.")
         return result
@@ -219,9 +255,9 @@ class CDFToAPI:
 
         exporter = exporters.DMSExporter(existing=existing, export_components=components, drop_data=drop_data)
 
-        if not self._client:
+        if not self._state.client:
             raise NeatSessionError("No client provided!")
 
-        result = self._state.rule_store.export_to_cdf(exporter, self._client, dry_run)
+        result = self._state.rule_store.export_to_cdf(exporter, self._state.client, dry_run)
         print("You can inspect the details with the .inspect.outcome.data_model(...) method.")
         return result

cognite/neat/_session/exceptions.py

@@ -2,17 +2,21 @@ import functools
 from collections.abc import Callable
 from typing import Any
 
-from cognite.neat._issues.errors import CDFMissingClientError
+from cognite.neat._issues.errors import CDFMissingClientError, NeatImportError
 
 from ._collector import _COLLECTOR
 
 try:
     from rich import print
+    from rich.markup import escape
 
     _PREFIX = "[bold red][ERROR][/bold red]"
 except ImportError:
     _PREFIX = "[ERROR]"
 
+    def escape(x: Any, *_: Any, **__: Any) -> Any:  # type: ignore[misc]
+        return x
+
 
 class NeatSessionError(Exception):
     """Base class for all exceptions raised by the NeatSession class."""
@@ -29,8 +33,8 @@ def _session_method_wrapper(func: Callable, cls_name: str):
         except NeatSessionError as e:
             action = _get_action()
             print(f"{_PREFIX} Cannot {action}: {e}")
-        except CDFMissingClientError as e:
-            print(f"{_PREFIX} {e.as_message()}")
+        except (CDFMissingClientError, NeatImportError) as e:
+            print(f"{_PREFIX} {escape(e.as_message())}")
         except ModuleNotFoundError as e:
             if e.name == "neatengine":
                 action = _get_action()