cognite-neat 0.106.0__py3-none-any.whl → 0.108.0__py3-none-any.whl

cognite/neat/_session/_drop.py

@@ -31,7 +31,11 @@ class DropAPI:
             ```
         """
         type_list = type if isinstance(type, list) else [type]
-        uri_type_type = dict((v, k) for k, v in self._state.instances.store.queries.types.items())
+
+        # Temporary solution until we agree on the form of specifying named graphs
+        # it will default to the default named graph
+        named_graph = self._state.instances.store.default_named_graph
+        uri_type_type = dict((v, k) for k, v in self._state.instances.store.queries.types(named_graph).items())
         selected_uri_by_type: dict[URIRef, str] = {}
         for type_item in type_list:
             if type_item not in uri_type_type:

cognite/neat/_session/_inspect.py

@@ -94,7 +94,8 @@ class InspectIssues:
         if self._state.rule_store.provenance:
             issues = self._state.rule_store.last_issues
         elif self._state.instances.store.provenance:
-            issues = self._state.instances.store.provenance[-1].target_entity.issues
+            last_change = self._state.instances.store.provenance[-1]
+            issues = last_change.target_entity.issues
         else:
             self._print("No issues found.")
             return pd.DataFrame() if return_dataframe else None
@@ -231,7 +232,7 @@ class InspectUploadOutcome:
                                     if i < 50:
                                         lines.append(f"  * {v}")
                                     elif i == 50 and total > 50:
-                                        lines.append(f"  * ... {total-50} more")
+                                        lines.append(f"  * ... {total - 50} more")
                                     elif i == 50 and total == 50:
                                         lines.append(f"  * {v}")
                                 else:

cognite/neat/_session/_mapping.py

@@ -1,7 +1,13 @@
 from cognite.neat._issues import IssueList
 from cognite.neat._rules.models import DMSRules
 from cognite.neat._rules.models.mapping import load_classic_to_core_mapping
-from cognite.neat._rules.transformers import AsParentPropertyId, IncludeReferenced, RuleMapper
+from cognite.neat._rules.transformers import (
+    AsParentPropertyId,
+    ChangeViewPrefix,
+    IncludeReferenced,
+    RuleMapper,
+    RulesTransformer,
+)
 
 from ._state import SessionState
 from .exceptions import NeatSessionError, session_class_wrapper
@@ -18,7 +24,7 @@ class DataModelMappingAPI:
     def __init__(self, state: SessionState):
         self._state = state
 
-    def classic_to_core(self, company_prefix: str, use_parent_property_name: bool = True) -> IssueList:
+    def classic_to_core(self, company_prefix: str | None = None, use_parent_property_name: bool = True) -> IssueList:
         """Map classic types to core types.
 
         Note this automatically creates an extended CogniteCore model.
@@ -45,10 +51,15 @@ class DataModelMappingAPI:
         if self._state.client is None:
             raise NeatSessionError("Client is required to map classic to core")
 
-        transformers = [
-            RuleMapper(load_classic_to_core_mapping(company_prefix, rules.metadata.space, rules.metadata.version)),
-            IncludeReferenced(self._state.client),
-        ]
+        transformers: list[RulesTransformer] = []
+        if company_prefix:
+            transformers.append(ChangeViewPrefix("Classic", company_prefix))
+        transformers.extend(
+            [
+                RuleMapper(load_classic_to_core_mapping(company_prefix, rules.metadata.space, rules.metadata.version)),
+                IncludeReferenced(self._state.client),
+            ]
+        )
         if use_parent_property_name:
             transformers.append(AsParentPropertyId(self._state.client))
         return self._state.rule_transform(*transformers)

cognite/neat/_session/_prepare.py

@@ -7,7 +7,6 @@ from rdflib import URIRef
 from cognite.neat._constants import (
     get_default_prefixes_and_namespaces,
 )
-from cognite.neat._graph import extractors
 from cognite.neat._graph.transformers import (
     AttachPropertyFromTargetToSource,
     ConnectionToLiteral,
@@ -24,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,
@@ -348,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:
@@ -545,13 +508,3 @@ class DataModelPrepareAPI:
                 "Please set the client in the session, NeatSession(client=client)."
             )
         return self._state.rule_transform(IncludeReferenced(self._state.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))

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
 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
@@ -32,6 +36,22 @@ class ReadAPI:
         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:
@@ -78,6 +98,45 @@ 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,
+        skip_cognite_views: bool = True,
+    ) -> 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.
+            skip_cognite_views: If True, all Cognite Views are skipped. For example, if you have the CogniteAsset
+                view in you data model, it will ont be used to extract instances.
+
+        Returns:
+            IssueList: A list of issues that occurred during the extraction.
+
+        """
+        return self._graph(data_model_id, instance_space, skip_cognite_views, unpack_json=False)
+
+    def _graph(
+        self,
+        data_model_id: DataModelIdentifier,
+        instance_space: str | SequenceNotStr[str] | None = None,
+        skip_cognite_views: bool = True,
+        unpack_json: bool = False,
+        str_to_ideal_type: bool = False,
+    ) -> IssueList:
+        extractor = extractors.DMSGraphExtractor.from_data_model_id(
+            # We are skipping the Cognite Views
+            data_model_id,
+            self._get_client,
+            instance_space=instance_space,
+            skip_cognite_views=skip_cognite_views,
+            unpack_json=unpack_json,
+            str_to_ideal_type=str_to_ideal_type,
+        )
+        return self._state.write_graph(extractor)
+
 
 @session_class_wrapper
 class CDFClassicAPI(BaseReadAPI):
@@ -92,7 +151,12 @@ class CDFClassicAPI(BaseReadAPI):
             raise ValueError("No client provided. Please provide a client to read a data model.")
         return self._state.client
 
-    def graph(self, root_asset_external_id: str, limit_per_type: int | None = None) -> IssueList:
+    def graph(
+        self,
+        root_asset_external_id: str,
+        limit_per_type: int | None = None,
+        identifier: Literal["id", "externalId"] = "id",
+    ) -> IssueList:
         """Reads the classic knowledge graph from CDF.
 
         The Classic Graph consists of the following core resource type.
@@ -127,6 +191,8 @@ class CDFClassicAPI(BaseReadAPI):
         Args:
             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.
+            identifier: The identifier to use for the core nodes. Note selecting "id" can cause issues if the external
+                ID of the core nodes is missing. Default is "id".
 
         Returns:
             IssueList: A list of issues that occurred during the extraction.
@@ -135,17 +201,58 @@ class CDFClassicAPI(BaseReadAPI):
             ```python
             neat.read.cdf.graph("root_asset_external_id")
             ```
-
         """
-        extractor = extractors.ClassicGraphExtractor(
-            self._get_client, root_asset_external_id=root_asset_external_id, limit_per_type=limit_per_type
+        return self._graph(
+            root_asset_external_id, limit_per_type, identifier, reference_timeseries=False, reference_files=False
         )
 
-        issues = self._state.instances.store.write(extractor)
-        issues.action = "Read Classic Graph"
-        if issues:
+    def _graph(
+        self,
+        root_asset_external_id: str,
+        limit_per_type: int | None = None,
+        identifier: Literal["id", "externalId"] = "id",
+        reference_timeseries: bool = False,
+        reference_files: bool = False,
+    ) -> IssueList:
+        namespace = CLASSIC_CDF_NAMESPACE
+        extractor = extractors.ClassicGraphExtractor(
+            self._get_client,
+            root_asset_external_id=root_asset_external_id,
+            limit_per_type=limit_per_type,
+            namespace=namespace,
+            prefix="Classic",
+            identifier=identifier,
+        )
+        extract_issues = self._state.write_graph(extractor)
+        if identifier == "externalId":
+            self._state.quoted_source_identifiers = True
+
+        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.
+        prepare_issues = self._state.rule_store.transform(
+            ClassicPrepareCore(namespace, reference_timeseries, reference_files)
+        )
+        # Update the instance store with the latest rules
+        information_rules = self._state.rule_store.last_verified_information_rules
+        self._state.instances.store.rules[self._state.instances.store.default_named_graph] = information_rules
+
+        all_issues = IssueList(extract_issues + prepare_issues)
+        # Update the provenance with all issue.
+        object.__setattr__(self._state.instances.store.provenance[-1].target_entity, "issues", all_issues)
+        all_issues.action = "Read Classic Graph"
+        if all_issues:
             print("Use the .inspect.issues() for more details")
-        return issues
+
+        return all_issues
 
 
 @session_class_wrapper
@@ -182,7 +289,6 @@ class ExcelReadAPI(BaseReadAPI):
 class ExcelExampleAPI(BaseReadAPI):
     """Used as example for reading some data model into the NeatSession."""
 
-    @property
     def pump_example(self) -> IssueList:
         """Reads the Hello World pump example into the NeatSession."""
         importer: importers.ExcelImporter = importers.ExcelImporter(catalog.hello_world_pump)
@@ -396,7 +502,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

@@ -3,9 +3,12 @@ from cognite.client import data_modeling as dm
 
 from cognite.neat._client import NeatClient
 from cognite.neat._constants import COGNITE_MODELS
+from cognite.neat._graph.transformers import SetNeatType
 from cognite.neat._issues import IssueList
+from cognite.neat._issues.errors import NeatValueError
 from cognite.neat._rules.models import DMSRules
 from cognite.neat._rules.transformers import SetIDDMSModel
+from cognite.neat._utils.text import humanize_collection
 
 from ._state import SessionState
 from .exceptions import NeatSessionError, session_class_wrapper
@@ -44,3 +47,27 @@ class SetAPI:
         if self._verbose:
             print(f"Client set to {self._state.client.config.project} CDF project.")
         return None
+
+    def _instance_sub_type(self, type: str, property: str, drop_property: bool = False) -> None:
+        """Sets the sub type of an instance based on the property."""
+        type_uri = self._state.instances.store.queries.type_uri(type)
+        property_uri = self._state.instances.store.queries.property_uri(property)
+
+        if not type_uri:
+            raise NeatValueError(f"Type {type} does not exist in the graph.")
+        elif len(type_uri) > 1:
+            raise NeatValueError(f"{type} has multiple ids found in the graph: {humanize_collection(type_uri)}.")
+
+        if not property_uri:
+            raise NeatValueError(f"Property {property} does not exist in the graph.")
+        elif len(type_uri) > 1:
+            raise NeatValueError(
+                f"{property} has multiple ids found in the graph: {humanize_collection(property_uri)}."
+            )
+
+        if not self._state.instances.store.queries.type_with_property(type_uri[0], property_uri[0]):
+            raise NeatValueError(f"Property {property} is not defined for type {type}.")
+
+        self._state.instances.store.transform(SetNeatType(type_uri[0], property_uri[0], drop_property))
+
+        return None

cognite/neat/_session/_show.py

@@ -112,7 +112,7 @@ class ShowDataModelAPI(ShowBaseAPI):
         else:
             # This should never happen, but we need to handle it to satisfy mypy
             raise NeatSessionError(
-                f"Unsupported type {type(rules) }. Make sure you have either information or DMS rules."
+                f"Unsupported type {type(rules)}. Make sure you have either information or DMS rules."
             )
         identifier = to_directory_compatible(str(rules.metadata.identifier))
         name = f"{identifier}.html"
@@ -201,7 +201,7 @@ class ShowDataModelImplementsAPI(ShowBaseAPI):
         else:
             # This should never happen, but we need to handle it to satisfy mypy
             raise NeatSessionError(
-                f"Unsupported type {type(rules) }. Make sure you have either information or DMS rules."
+                f"Unsupported type {type(rules)}. Make sure you have either information or DMS rules."
             )
         identifier = to_directory_compatible(str(rules.metadata.identifier))
         name = f"{identifier}_implements.html"
@@ -364,7 +364,7 @@ class ShowInstanceAPI(ShowBaseAPI):
                 'Try setting [bold]NeatSession(storage="oxigraph")[/bold] enable Oxigraph store.'
             )
 
-        if not self._state.instances.store.graph:
+        if not self._state.instances.store.dataset:
             raise NeatSessionError("No instances available. Try using [bold].read[/bold] to load instances.")
 
         di_graph = self._generate_instance_di_graph_and_types()
@@ -395,7 +395,7 @@ class ShowInstanceAPI(ShowBaseAPI):
             object,
             subject_type,
             object_type,
-        ) in self._state.instances.store.graph.query(query):
+        ) in self._state.instances.store.dataset.query(query):
             subject = remove_namespace_from_uri(subject)
             property_ = remove_namespace_from_uri(property_)
             object = remove_namespace_from_uri(object)

cognite/neat/_session/_state.py

@@ -2,6 +2,7 @@ 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
@@ -21,11 +22,14 @@ class SessionState:
         self.rule_store = NeatRulesStore()
         self.last_reference: DMSRules | InformationRules | None = None
         self.client = client
+        self.quoted_source_identifiers = False
 
     def rule_transform(self, *transformer: RulesTransformer) -> IssueList:
         if not transformer:
             raise NeatSessionError("No transformers provided.")
         first_transformer = transformer[0]
+
+        # This should not be allowed to be done automatically
         pruned = self.rule_store.prune_until_compatible(first_transformer)
         if pruned:
             type_hint = first_transformer.transform_type_hint()
@@ -62,6 +66,13 @@ class SessionState:
             issues.hint = "Use the .inspect.issues() for more details."
         return issues
 
+    def write_graph(self, extractor: KnowledgeGraphExtractor) -> IssueList:
+        extract_issues = self.instances.store.write(extractor)
+        issues = self.rule_store.import_graph(extractor)
+        self.instances.store.add_rules(self.rule_store.last_verified_information_rules)
+        issues.extend(extract_issues)
+        return issues
+
 
 @dataclass
 class InstancesState:
@@ -74,7 +85,7 @@ class InstancesState:
     def store(self) -> NeatGraphStore:
         if not self.has_store:
             if self.store_type == "oxigraph":
-                self._store = NeatGraphStore.from_oxi_store()
+                self._store = NeatGraphStore.from_oxi_local_store()
             else:
                 self._store = NeatGraphStore.from_memory_store()
         return cast(NeatGraphStore, self._store)

cognite/neat/_session/_to.py

@@ -1,3 +1,5 @@
+import warnings
+import zipfile
 from collections.abc import Collection
 from pathlib import Path
 from typing import Any, Literal, overload
@@ -80,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.trig",
+                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
@@ -182,6 +220,9 @@ class CDFToAPI:
             self._state.instances.store,
             instance_space=space,
             client=client,
+            # In case urllib.parse.quote() was run on the extraction, we need to run
+            # urllib.parse.unquote() on the load.
+            unquote_external_ids=self._state.quoted_source_identifiers,
         )
         result = loader.load_into_cdf(client)
         self._state.instances.outcome.append(result)

cognite/neat/_session/_wizard.py

@@ -26,7 +26,7 @@ _T_Option = TypeVar("_T_Option")
 
 
 def _selection(message: str, options: Sequence[_T_Option]) -> _T_Option:
-    option_text = "\n  ".join([f"{i+1}) {option}" for i, option in enumerate(options)])
+    option_text = "\n  ".join([f"{i + 1}) {option}" for i, option in enumerate(options)])
     selected_index = (
         IntPrompt().ask(f"{message}\n  {option_text}\n", choices=list(map(str, range(1, len(options) + 1)))) - 1
     )

cognite/neat/_session/exceptions.py

@@ -2,17 +2,22 @@ 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 cognite.neat._issues.errors._general import NeatValueError
 
 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 +34,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, NeatValueError) as e:
+            print(f"{_PREFIX} {escape(e.as_message())}")
         except ModuleNotFoundError as e:
             if e.name == "neatengine":
                 action = _get_action()