cognite-neat 0.109.3__py3-none-any.whl → 0.110.0__py3-none-any.whl

cognite/neat/_session/_show.py

@@ -1,16 +1,13 @@
 import colorsys
 import random
-from typing import Any, cast
+from typing import Any
 
 import networkx as nx
 from IPython.display import HTML, display
 from pyvis.network import Network as PyVisNetwork  # type: ignore
 
 from cognite.neat._constants import IN_NOTEBOOK, IN_PYODIDE
-from cognite.neat._rules._constants import EntityTypes
-from cognite.neat._rules.models.dms._rules import DMSRules
-from cognite.neat._rules.models.entities._single_value import ClassEntity, ViewEntity
-from cognite.neat._rules.models.information._rules import InformationRules
+from cognite.neat._rules.analysis._base import RulesAnalysis
 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, uri_display_name
@@ -99,79 +96,21 @@ class ShowDataModelAPI(ShowBaseAPI):
 
     def __call__(self) -> Any:
         if self._state.rule_store.empty:
-            raise NeatSessionError("No data model available. Try using [bold].read[/bold] to read a new data model.")
+            raise NeatSessionError("No data model available. Try using [bold].read[/bold] to read a data model.")
+
         last_target = self._state.rule_store.provenance[-1].target_entity
         rules = last_target.dms or last_target.information
+        analysis = RulesAnalysis(dms=last_target.dms, information=last_target.information)
 
         if last_target.dms is not None:
-            di_graph = self._generate_dms_di_graph(last_target.dms)
+            di_graph = analysis._dms_di_graph(format="data-model")
         else:
-            di_graph = self._generate_info_di_graph(last_target.information)
+            di_graph = analysis._info_di_graph(format="data-model")
+
         identifier = to_directory_compatible(str(rules.metadata.identifier))
         name = f"{identifier}.html"
-
         return self._generate_visualization(di_graph, name)
 
-    def _generate_dms_di_graph(self, rules: DMSRules) -> nx.DiGraph:
-        """Generate a DiGraph from the last verified DMS rules."""
-        di_graph = nx.DiGraph()
-
-        # Views with properties or used as ValueType
-        # If a view is not used in properties or as ValueType, it is not added to the graph
-        # as we typically do not have the properties for it.
-        used_views = {prop_.view for prop_ in rules.properties} | {
-            prop_.value_type for prop_ in rules.properties if isinstance(prop_.value_type, ViewEntity)
-        }
-
-        # Add nodes and edges from Views sheet
-        for view in rules.views:
-            if view.view not in used_views:
-                continue
-            # if possible use humanreadable label coming from the view name
-            if not di_graph.has_node(view.view.suffix):
-                di_graph.add_node(view.view.suffix, label=view.view.suffix)
-
-        # Add nodes and edges from Properties sheet
-        for prop_ in rules.properties:
-            if prop_.connection and isinstance(prop_.value_type, ViewEntity):
-                if not di_graph.has_node(prop_.view.suffix):
-                    di_graph.add_node(prop_.view.suffix, label=prop_.view.suffix)
-                di_graph.add_edge(
-                    prop_.view.suffix,
-                    prop_.value_type.suffix,
-                    label=prop_.name or prop_.view_property,
-                )
-
-        return di_graph
-
-    def _generate_info_di_graph(self, rules: InformationRules) -> nx.DiGraph:
-        """Generate DiGraph representing information data model."""
-
-        di_graph = nx.DiGraph()
-
-        # Add nodes and edges from Views sheet
-        for class_ in rules.classes:
-            # if possible use humanreadable label coming from the view name
-            if not di_graph.has_node(class_.class_.suffix):
-                di_graph.add_node(
-                    class_.class_.suffix,
-                    label=class_.name or class_.class_.suffix,
-                )
-
-        # Add nodes and edges from Properties sheet
-        for prop_ in rules.properties:
-            if prop_.type_ == EntityTypes.object_property:
-                if not di_graph.has_node(prop_.class_.suffix):
-                    di_graph.add_node(prop_.class_.suffix, label=prop_.class_.suffix)
-
-                di_graph.add_edge(
-                    prop_.class_.suffix,
-                    cast(ClassEntity, prop_.value_type).suffix,
-                    label=prop_.name or prop_.property_,
-                )
-
-        return di_graph
-
 
 @session_class_wrapper
 class ShowDataModelImplementsAPI(ShowBaseAPI):
@@ -185,67 +124,16 @@ class ShowDataModelImplementsAPI(ShowBaseAPI):
 
         last_target = self._state.rule_store.provenance[-1].target_entity
         rules = last_target.dms or last_target.information
+        analysis = RulesAnalysis(dms=last_target.dms, information=last_target.information)
 
         if last_target.dms is not None:
-            di_graph = self._generate_dms_di_graph(last_target.dms)
+            di_graph = analysis._dms_di_graph(format="implements")
         else:
-            di_graph = self._generate_info_di_graph(last_target.information)
+            di_graph = analysis._info_di_graph(format="implements")
         identifier = to_directory_compatible(str(rules.metadata.identifier))
         name = f"{identifier}_implements.html"
         return self._generate_visualization(di_graph, name)
 
-    def _generate_dms_di_graph(self, rules: DMSRules) -> nx.DiGraph:
-        """Generate a DiGraph from the last verified DMS rules."""
-        di_graph = nx.DiGraph()
-
-        # Add nodes and edges from Views sheet
-        for view in rules.views:
-            # add implements as edges
-            if view.implements:
-                if not di_graph.has_node(view.view.suffix):
-                    di_graph.add_node(view.view.suffix, label=view.view.suffix)
-                for implement in view.implements:
-                    if not di_graph.has_node(implement.suffix):
-                        di_graph.add_node(implement.suffix, label=implement.suffix)
-
-                    di_graph.add_edge(
-                        view.view.suffix,
-                        implement.suffix,
-                        label="implements",
-                        dashes=True,
-                    )
-
-        return di_graph
-
-    def _generate_info_di_graph(self, rules: InformationRules) -> nx.DiGraph:
-        """Generate DiGraph representing information data model."""
-
-        di_graph = nx.DiGraph()
-
-        # Add nodes and edges from Views sheet
-        for class_ in rules.classes:
-            # if possible use human readable label coming from the view name
-
-            # add subClassOff as edges
-            if class_.implements:
-                if not di_graph.has_node(class_.class_.suffix):
-                    di_graph.add_node(
-                        class_.class_.suffix,
-                        label=class_.name or class_.class_.suffix,
-                    )
-
-                for parent in class_.implements:
-                    if not di_graph.has_node(parent.suffix):
-                        di_graph.add_node(parent.suffix, label=parent.suffix)
-                    di_graph.add_edge(
-                        class_.class_.suffix,
-                        parent.suffix,
-                        label="implements",
-                        dashes=True,
-                    )
-
-        return di_graph
-
 
 @session_class_wrapper
 class ShowDataModelProvenanceAPI(ShowBaseAPI):

cognite/neat/_session/_state.py

@@ -37,7 +37,6 @@ class SessionState:
         last_entity = self.rule_store.provenance[-1].target_entity
         issues.action = f"{start} → {last_entity.display_name}"
         issues.hint = "Use the .inspect.issues() for more details."
-        self.instances.store.add_rules(last_entity.information)
         return issues
 
     def rule_import(self, importer: BaseImporter, enable_manual_edit: bool = False) -> IssueList:
@@ -61,7 +60,6 @@ class SessionState:
     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
 

cognite/neat/_session/_subset.py

@@ -0,0 +1,64 @@
+import warnings
+
+from cognite.neat._alpha import AlphaFlags
+from cognite.neat._issues._base import IssueList
+from cognite.neat._rules.models.entities._single_value import ClassEntity, ViewEntity
+from cognite.neat._rules.transformers import SubsetDMSRules, SubsetInformationRules
+
+from ._state import SessionState
+from .exceptions import NeatSessionError, session_class_wrapper
+
+try:
+    from rich import print
+except ImportError:
+    ...
+
+
+@session_class_wrapper
+class SubsetAPI:
+    """
+    Subset data model and instances in the session based on the desired subset of concepts.
+
+    """
+
+    def __init__(self, state: SessionState):
+        self._state = state
+
+    def data_model(self, concepts: str | list[str]) -> IssueList:
+        if self._state.rule_store.empty:
+            raise NeatSessionError("No rules to set the data model ID.")
+
+        warnings.filterwarnings("default")
+        AlphaFlags.data_model_subsetting.warn()
+
+        dms = self._state.rule_store.provenance[-1].target_entity.dms
+        information = self._state.rule_store.provenance[-1].target_entity.information
+
+        if dms:
+            views = {
+                ViewEntity(
+                    space=dms.metadata.space,
+                    externalId=concept,
+                    version=dms.metadata.version,
+                )
+                for concept in concepts
+            }
+
+            issues = self._state.rule_transform(SubsetDMSRules(views=views))
+            if not issues:
+                after = len(self._state.rule_store.last_verified_dms_rules.views)
+
+        elif information:
+            classes = {ClassEntity(prefix=information.metadata.space, suffix=concept) for concept in concepts}
+
+            issues = self._state.rule_transform(SubsetInformationRules(classes=classes))
+            if not issues:
+                after = len(self._state.rule_store.last_verified_information_rules.classes)
+
+        else:
+            raise NeatSessionError("Something went terrible wrong. Please contact the neat team.")
+
+        if not issues:
+            print(f"Subset to {after} concepts.")
+
+        return issues

cognite/neat/_session/_to.py

@@ -5,14 +5,18 @@ from pathlib import Path
 from typing import Any, Literal, overload
 
 from cognite.client import data_modeling as dm
+from cognite.client.data_classes.data_modeling import DataModelIdentifier
 
 from cognite.neat._alpha import AlphaFlags
 from cognite.neat._constants import COGNITE_MODELS
 from cognite.neat._graph import loaders
+from cognite.neat._issues import IssueList, catch_issues
 from cognite.neat._rules import exporters
 from cognite.neat._rules._constants import PATTERNS
 from cognite.neat._rules._shared import VerifiedRules
 from cognite.neat._rules.exporters._rules2dms import Component
+from cognite.neat._rules.importers import DMSImporter
+from cognite.neat._rules.models import DMSRules, InformationRules
 from cognite.neat._rules.models.dms import DMSMetadata
 from cognite.neat._utils.upload import UploadResultList
 
@@ -35,10 +39,10 @@ class ToAPI:
     def excel(
         self,
         io: Any,
-        include_reference: bool = True,
+        include_reference: bool | DataModelIdentifier = True,
         include_properties: Literal["same-space", "all"] = "all",
         add_empty_rows: bool = False,
-    ) -> None:
+    ) -> IssueList | None:
         """Export the verified data model to Excel.
 
         Args:
@@ -46,6 +50,7 @@ class ToAPI:
             include_reference: If True, the reference data model will be included. Defaults to True.
                 Note that this only applies if you have created the data model using the
                 create.enterprise_model(...), create.solution_model(), or create.data_product_model() methods.
+                You can also provide a DataModelIdentifier directly, which will be read from CDF
             include_properties: The properties to include in the Excel file. Defaults to "all".
                 - "same-space": Only properties that are in the same space as the data model will be included.
             add_empty_rows: If True, empty rows will be added between each component. Defaults to False.
@@ -71,19 +76,45 @@ class ToAPI:
             dms_rules_file_name = "dms_rules.xlsx"
             neat.to.excel(dms_rules_file_name, include_reference=True)
             ```
+
+        Example:
+            Read the data model ("my_space", "ISA95Model", "v5") and export it to an excel file with the
+            CogniteCore model in the reference sheets.
+            ```python
+            client = CogniteClient()
+            neat = NeatSession(client)
+
+            neat.read.cdf(("my_space", "ISA95Model", "v5"))
+            dms_rules_file_name = "dms_rules.xlsx"
+            neat.to.excel(dms_rules_file_name, include_reference=("cdf_cdm", "CogniteCore", "v1"))
         """
         reference_rules_with_prefix: tuple[VerifiedRules, str] | None = None
         include_properties = include_properties.strip().lower()
 
-        if include_reference and self._state.last_reference:
-            if (
-                isinstance(self._state.last_reference.metadata, DMSMetadata)
-                and self._state.last_reference.metadata.as_data_model_id() in COGNITE_MODELS
-            ):
-                prefix = "CDM"
+        if include_reference is not False:
+            if include_reference is True and self._state.last_reference is not None:
+                ref_rules: InformationRules | DMSRules | None = self._state.last_reference
+            elif include_reference is True:
+                ref_rules = None
             else:
+                if not self._state.client:
+                    raise NeatSessionError("No client provided!")
+                ref_rules = None
+                with catch_issues() as issues:
+                    ref_read = DMSImporter.from_data_model_id(self._state.client, include_reference).to_rules()
+                    if ref_read.rules is not None:
+                        ref_rules = ref_read.rules.as_verified_rules()
+                if ref_rules is None or issues.has_errors:
+                    issues.action = f"Read {include_reference}"
+                    return issues
+            if ref_rules is not None:
                 prefix = "Ref"
-            reference_rules_with_prefix = self._state.last_reference, prefix
+                if (
+                    isinstance(ref_rules.metadata, DMSMetadata)
+                    and ref_rules.metadata.as_data_model_id() in COGNITE_MODELS
+                ):
+                    prefix = "CDM"
+                reference_rules_with_prefix = ref_rules, prefix
 
         if include_properties == "same-space":
             warnings.filterwarnings("default")
@@ -95,7 +126,8 @@ class ToAPI:
             add_empty_rows=add_empty_rows,
             include_properties=include_properties,  # type: ignore
         )
-        return self._state.rule_store.export_to_file(exporter, Path(io))
+        self._state.rule_store.export_to_file(exporter, Path(io))
+        return None
 
     def session(self, io: Any) -> None:
         """Export the current session to a file.
@@ -212,12 +244,30 @@ class CDFToAPI:
     def instances(
         self,
         space: str | None = None,
+        space_property: str | None = None,
     ) -> UploadResultList:
         """Export the verified DMS instances to CDF.
 
         Args:
             space: Name of instance space to use. Default is to suffix the schema space with '_instances'.
-            Note this space is required to be different than the space with the data model.
+                Note this space is required to be different from the space with the data model.
+            space_property: This is an alternative to the 'space' argument. If provided, the space will set to the
+                value of the property with the given name for each instance. If the property is not found, the
+                'space' argument will be used. Defaults to None.
+
+        Returns:
+            UploadResultList: The result of the upload.
+
+        Example:
+            Export instances to CDF
+            ```python
+            neat.to.cdf.instances()
+            ```
+
+            Export instances to CDF using the `dataSetId` property as the space
+            ```python
+            neat.to.cdf.instances(space_property="dataSetId")
+            ```
 
         """
         if not self._state.client:
@@ -233,8 +283,9 @@ class CDFToAPI:
         if not client.data_modeling.spaces.retrieve(space):
             client.data_modeling.spaces.apply(dm.SpaceApply(space=space))
 
-        loader = loaders.DMSLoader.from_rules(
+        loader = loaders.DMSLoader(
             self._state.rule_store.last_verified_dms_rules,
+            self._state.rule_store.last_verified_information_rules,
             self._state.instances.store,
             instance_space=space,
             client=client,