cognite-neat 0.109.4__py3-none-any.whl → 0.111.0__py3-none-any.whl

cognite/neat/_session/_prepare.py

@@ -1,3 +1,4 @@
+import warnings
 from collections.abc import Callable
 from typing import Any
 
@@ -14,6 +15,7 @@ from cognite.neat._graph.transformers._rdfpath import MakeConnectionOnExactMatch
 from cognite.neat._issues import IssueList
 from cognite.neat._issues.errors import NeatValueError
 from cognite.neat._rules.transformers import PrefixEntities, StandardizeNaming
+from cognite.neat._rules.transformers._converters import StandardizeSpaceAndVersion
 from cognite.neat._utils.text import humanize_collection
 
 from ._state import SessionState
@@ -270,5 +272,15 @@ class DataModelPrepareAPI:
         For classes/views/containers, the naming will be standardized to PascalCase.
         For properties, the naming will be standardized to camelCase.
         """
+        warnings.filterwarnings("default")
         AlphaFlags.standardize_naming.warn()
         return self._state.rule_transform(StandardizeNaming())
+
+    def standardize_space_and_version(self) -> IssueList:
+        """Standardize space and version in the data model.
+
+        This method will standardize the space and version in the data model to the Cognite standard.
+        """
+        warnings.filterwarnings("default")
+        AlphaFlags.standardize_space_and_version.warn()
+        return self._state.rule_transform(StandardizeSpaceAndVersion())

cognite/neat/_session/_read.py

@@ -50,6 +50,7 @@ class ReadAPI:
         self.csv = CSVReadAPI(state, verbose)
         self.yaml = YamlReadAPI(state, verbose)
         self.xml = XMLReadAPI(state, verbose)
+        self.examples = Examples(state)
 
     def session(self, io: Any) -> None:
         """Reads a Neat Session from a zip file.
@@ -152,6 +153,45 @@ class CDFReadAPI(BaseReadAPI):
         )
         return self._state.write_graph(extractor)
 
+    def raw(
+        self,
+        db_name: str,
+        table_name: str,
+        type: str | None = None,
+        foreign_keys: str | SequenceNotStr[str] | None = None,
+        unpack_json: bool = False,
+        str_to_ideal_type: bool = False,
+    ) -> IssueList:
+        """Reads a raw table from CDF to the knowledge graph.
+
+        Args:
+            db_name: The name of the database
+            table_name: The name of the table, this will be assumed to be the type of the instances.
+            type: The type of instances in the table. If None, the table name will be used.
+            foreign_keys: The name of the columns that are foreign keys. If None, no foreign keys are used.
+            unpack_json: If True, the JSON objects will be unpacked into the graph.
+            str_to_ideal_type: If True, the string values will be converted to ideal types.
+
+        Returns:
+            IssueList: A list of issues that occurred during the extraction.
+
+        Example:
+            ```python
+            neat.read.cdf.raw("my_db", "my_table", "Asset")
+            ```
+
+        """
+        extractor = extractors.RAWExtractor(
+            self._get_client,
+            db_name=db_name,
+            table_name=table_name,
+            table_type=type,
+            foreign_keys=foreign_keys,
+            unpack_json=unpack_json,
+            str_to_ideal_type=str_to_ideal_type,
+        )
+        return self._state.instances.store.write(extractor)
+
 
 @session_class_wrapper
 class CDFClassicAPI(BaseReadAPI):
@@ -228,6 +268,8 @@ class CDFClassicAPI(BaseReadAPI):
         identifier: Literal["id", "externalId"] = "id",
         reference_timeseries: bool = False,
         reference_files: bool = False,
+        unpack_metadata: bool = False,
+        skip_sequence_rows: bool = False,
     ) -> IssueList:
         namespace = CLASSIC_CDF_NAMESPACE
         extractor = extractors.ClassicGraphExtractor(
@@ -237,7 +279,11 @@ class CDFClassicAPI(BaseReadAPI):
             namespace=namespace,
             prefix="Classic",
             identifier=identifier,
+            unpack_metadata=unpack_metadata,
+            skip_sequence_rows=skip_sequence_rows,
         )
+        self._state.instances.neat_prefix_by_predicate_uri.update(extractor.neat_prefix_by_predicate_uri)
+        self._state.instances.neat_prefix_by_type_uri.update(extractor.neat_prefix_by_type_uri)
         extract_issues = self._state.write_graph(extractor)
         if identifier == "externalId":
             self._state.quoted_source_identifiers = True
@@ -256,9 +302,6 @@ class CDFClassicAPI(BaseReadAPI):
         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.
@@ -286,7 +329,6 @@ class ExcelReadAPI(BaseReadAPI):
 
     def __init__(self, state: SessionState, verbose: bool) -> None:
         super().__init__(state, verbose)
-        self.examples = ExcelExampleAPI(state, verbose)
 
     def __call__(self, io: Any, enable_manual_edit: bool = False) -> IssueList:
         """Reads a Neat Excel Rules sheet to the graph store. The rules sheet may stem from an Information architect,
@@ -310,16 +352,6 @@ class ExcelReadAPI(BaseReadAPI):
         return self._state.rule_import(importers.ExcelImporter(path), enable_manual_edit)
 
 
-@session_class_wrapper
-class ExcelExampleAPI(BaseReadAPI):
-    """Used as example for reading some data model into the NeatSession."""
-
-    def pump_example(self) -> IssueList:
-        """Reads the Hello World pump example into the NeatSession."""
-        importer: importers.ExcelImporter = importers.ExcelImporter(catalog.hello_world_pump)
-        return self._state.rule_import(importer)
-
-
 @session_class_wrapper
 class YamlReadAPI(BaseReadAPI):
     def __call__(self, io: Any, format: Literal["neat", "toolkit"] = "neat") -> IssueList:
@@ -374,6 +406,9 @@ class CSVReadAPI(BaseReadAPI):
     """
 
     def __call__(self, io: Any, type: str, primary_key: str) -> None:
+        warnings.filterwarnings("default")
+        AlphaFlags.csv_read.warn()
+
         engine = import_engine()
         engine.set.format = "csv"
         engine.set.file = NeatReader.create(io).materialize_path()
@@ -429,6 +464,9 @@ class XMLReadAPI(BaseReadAPI):
             - remove associations between nodes that do not exist in the extracted graph
             - remove edges to nodes that do not exist in the extracted graph
         """
+        warnings.filterwarnings("default")
+        AlphaFlags.dexpi_read.warn()
+
         path = NeatReader.create(io).materialize_path()
         engine = import_engine()
         engine.set.format = "dexpi"
@@ -480,6 +518,9 @@ class XMLReadAPI(BaseReadAPI):
             - remove unused attributes
             - remove edges to nodes that do not exist in the extracted graph
         """
+        warnings.filterwarnings("default")
+        AlphaFlags.aml_read.warn()
+
         path = NeatReader.create(io).materialize_path()
         engine = import_engine()
         engine.set.format = "aml"
@@ -519,7 +560,6 @@ class RDFReadAPI(BaseReadAPI):
 
     def __init__(self, state: SessionState, verbose: bool) -> None:
         super().__init__(state, verbose)
-        self.examples = RDFExamples(state)
 
     def ontology(self, io: Any) -> IssueList:
         """Reads an OWL ontology source into NeatSession.
@@ -532,6 +572,9 @@ class RDFReadAPI(BaseReadAPI):
             neat.read.rdf.ontology("url_or_path_to_owl_source")
             ```
         """
+        warnings.filterwarnings("default")
+        AlphaFlags.ontology_read.warn()
+
         reader = NeatReader.create(io)
         importer = importers.OWLImporter.from_file(reader.materialize_path(), source_name=f"file {reader!s}")
         return self._state.rule_import(importer)
@@ -547,10 +590,18 @@ class RDFReadAPI(BaseReadAPI):
             neat.read.rdf.imf("url_or_path_to_imf_source")
             ```
         """
+        warnings.filterwarnings("default")
+        AlphaFlags.imf_read.warn()
+
         reader = NeatReader.create(io)
         importer = importers.IMFImporter.from_file(reader.materialize_path(), source_name=f"file {reader!s}")
         return self._state.rule_import(importer)
 
+    def instances(self, io: Any) -> IssueList:
+        reader = NeatReader.create(io)
+        self._state.instances.store.write(extractors.RdfFileExtractor(reader.materialize_path()))
+        return IssueList()
+
     def __call__(
         self,
         io: Any,
@@ -574,21 +625,38 @@ class RDFReadAPI(BaseReadAPI):
                 raise ValueError(f"Expected ontology, imf types or instances, got {source}")
 
         elif type == "instances":
-            reader = NeatReader.create(io)
-            self._state.instances.store.write(extractors.RdfFileExtractor(reader.materialize_path()))
-            return IssueList()
+            return self.instances(io)
+
         else:
             raise NeatSessionError(f"Expected data model or instances, got {type}")
 
 
 @session_class_wrapper
-class RDFExamples:
-    """Used as example for reading some triples into the NeatSession knowledge grapgh."""
+class Examples:
+    """Used as example for reading various sources into NeatSession."""
 
     def __init__(self, state: SessionState) -> None:
         self._state = state
 
+    @property
+    def _get_client(self) -> NeatClient:
+        if self._state.client is None:
+            raise NeatValueError("No client provided. Please provide a client to read a data model.")
+        return self._state.client
+
     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))
         return IssueList()
+
+    def pump_example(self) -> IssueList:
+        """Reads the Hello World pump example into the NeatSession."""
+        importer: importers.ExcelImporter = importers.ExcelImporter(catalog.hello_world_pump)
+        return self._state.rule_import(importer)
+
+    def core_data_model(self) -> IssueList:
+        """Reads the core data model example into the NeatSession."""
+
+        cdm_v1 = DataModelId.load(("cdf_cdm", "CogniteCore", "v1"))
+        importer: importers.DMSImporter = importers.DMSImporter.from_data_model_id(self._get_client, cdm_v1)
+        return self._state.rule_import(importer)

cognite/neat/_session/_set.py

@@ -23,14 +23,20 @@ class SetAPI:
         self._verbose = verbose
         self.instances = SetInstances(state, verbose)
 
-    def data_model_id(self, new_model_id: dm.DataModelId | tuple[str, str, str]) -> IssueList:
+    def data_model_id(self, new_model_id: dm.DataModelId | tuple[str, str, str], name: str | None = None) -> 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>).
 
+        Args:
+            new_model_id (dm.DataModelId | tuple[str, str, str]): The new data model id.
+            name (str, optional): The display name of the data model. If not set, the external ID will be used
+                to generate the name.
+
         Example:
             Set a new data model id:
             ```python
             neat.set.data_model_id(("my_data_model_space", "My_Data_Model", "v1"))
+            neat.set.data_model_id(("my_data_model_space", "MyDataModel", "v1"), name="My Data Model")
             ```
         """
         if self._state.rule_store.empty:

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

@@ -1,6 +1,8 @@
 from pathlib import Path
 from typing import Literal, cast
 
+from rdflib import URIRef
+
 from cognite.neat._client import NeatClient
 from cognite.neat._graph.extractors import KnowledgeGraphExtractor
 from cognite.neat._issues import IssueList
@@ -37,7 +39,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 +62,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
 
@@ -76,6 +76,10 @@ class InstancesState:
         self.storage_path = storage_path
         self.issue_lists = IssueList()
         self.outcome = UploadResultList()
+        # These contain prefixes added by Neat at the extraction stage.
+        # We store them such that they can be removed in the load stage.
+        self.neat_prefix_by_predicate_uri: dict[URIRef, str] = {}
+        self.neat_prefix_by_type_uri: dict[URIRef, str] = {}
 
         # Ensure that error handling is done in the constructor
         self.store: NeatGraphStore = _session_method_wrapper(self._create_store, "NeatSession")()

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