cognite-neat 0.101.0__py3-none-any.whl → 0.103.0__py3-none-any.whl

cognite/neat/_session/_prepare.py

@@ -1,14 +1,26 @@
 import copy
-from collections.abc import Collection
+from collections.abc import Callable, Collection
 from datetime import datetime, timezone
-from typing import Literal, cast
+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 DEFAULT_NAMESPACE
-from cognite.neat._graph.transformers import RelationshipAsEdgeTransformer
+from cognite.neat._constants import (
+    DEFAULT_NAMESPACE,
+    get_default_prefixes_and_namespaces,
+)
+from cognite.neat._graph.transformers import (
+    AttachPropertyFromTargetToSource,
+    ConvertLiteral,
+    LiteralToEntity,
+    PruneDeadEndEdges,
+    PruneInstancesOfUnknownType,
+    PruneTypes,
+    RelationshipAsEdgeTransformer,
+    Transformers,
+)
 from cognite.neat._graph.transformers._rdfpath import MakeConnectionOnExactMatch
 from cognite.neat._rules._shared import InputRules, ReadRules
 from cognite.neat._rules.importers import DMSImporter
@@ -37,6 +49,10 @@ except ImportError:
 
 @session_class_wrapper
 class PrepareAPI:
+    """Apply various operations on the knowledge graph as a necessary preprocessing step before for instance
+    inferring a data model or exporting the knowledge graph to a desired destination.
+    """
+
     def __init__(self, client: NeatClient | None, state: SessionState, verbose: bool) -> None:
         self._state = state
         self._verbose = verbose
@@ -46,10 +62,97 @@ class PrepareAPI:
 
 @session_class_wrapper
 class InstancePrepareAPI:
+    """Operations to perform on instances of data in the knowledge graph."""
+
     def __init__(self, state: SessionState, verbose: bool) -> None:
         self._state = state
         self._verbose = verbose
 
+    def dexpi(self) -> None:
+        """Prepares extracted DEXPI graph for further usage in CDF
+
+        !!! note "This method bundles several graph transformers which"
+            - attach values of generic attributes to nodes
+            - create associations between nodes
+            - remove unused generic attributes
+            - remove associations between nodes that do not exist in the extracted graph
+            - remove edges to nodes that do not exist in the extracted graph
+
+        and therefore safeguard CDF from a bad graph
+
+        Example:
+            Apply Dexpi specific transformations:
+            ```python
+            neat.prepare.instances.dexpi()
+            ```
+        """
+
+        DEXPI = get_default_prefixes_and_namespaces()["dexpi"]
+
+        transformers = [
+            # Remove any instance which type is unknown
+            PruneInstancesOfUnknownType(),
+            # Directly connect generic attributes
+            AttachPropertyFromTargetToSource(
+                target_property=DEXPI.Value,
+                target_property_holding_new_property=DEXPI.Name,
+                target_node_type=DEXPI.GenericAttribute,
+                delete_target_node=True,
+            ),
+            # Directly connect associations
+            AttachPropertyFromTargetToSource(
+                target_property=DEXPI.ItemID,
+                target_property_holding_new_property=DEXPI.Type,
+                target_node_type=DEXPI.Association,
+                delete_target_node=True,
+            ),
+            # Remove unused generic attributes and associations
+            PruneTypes([DEXPI.GenericAttribute, DEXPI.Association]),
+            # Remove edges to nodes that do not exist in the extracted graph
+            PruneDeadEndEdges(),
+        ]
+
+        for transformer in transformers:
+            self._state.instances.store.transform(cast(Transformers, transformer))
+
+    def aml(self) -> None:
+        """Prepares extracted AutomationML graph for further usage in CDF
+
+        !!! note "This method bundles several graph transformers which"
+            - attach values of attributes to nodes
+            - remove unused attributes
+            - remove edges to nodes that do not exist in the extracted graph
+
+        and therefore safeguard CDF from a bad graph
+
+        Example:
+            Apply AML specific transformations:
+            ```python
+            neat.prepare.instances.aml()
+            ```
+        """
+
+        AML = get_default_prefixes_and_namespaces()["aml"]
+
+        transformers = [
+            # Remove any instance which type is unknown
+            PruneInstancesOfUnknownType(),
+            # Directly connect generic attributes
+            AttachPropertyFromTargetToSource(
+                target_property=AML.Value,
+                target_property_holding_new_property=AML.Name,
+                target_node_type=AML.Attribute,
+                delete_target_node=True,
+            ),
+            # Prune unused attributes
+            PruneTypes([AML.Attribute]),
+            # # Remove edges to nodes that do not exist in the extracted graph
+            PruneDeadEndEdges(),
+        ]
+
+        for transformer in transformers:
+            self._state.instances.store.transform(cast(Transformers, transformer))
+
     def make_connection_on_exact_match(
         self,
         source: tuple[str, str],
@@ -79,7 +182,32 @@ class InstancePrepareAPI:
             target value, as follows:
             (SourceType)-[connection]->(TargetType)
 
+        Example:
+            Make connection on exact match:
+            ```python
+            # From an active NeatSession
+            neat.read.csv("workitem.Table.csv",
+                          type = "Activity",
+                          primary_key="sourceId")
 
+            neat.read.csv("assets.Table.csv",
+                          type="Asset",
+                          primary_key="WMT_TAG_GLOBALID")
+
+            # Here we specify what column from the source table we should use when we link it with a column in the
+            # target table. In this case, it is the "workorderItemname" column in the source table
+            source = ("Activity", "workorderItemname")
+
+            # Here we give a name to the new property that is created when a match between the source and target is
+            # found
+            connection = "asset"
+
+            # Here we specify what column from the target table we should use when searching for a match.
+            # In this case, it is the "wmtTagName" column in the target table
+            target = ("Asset", "wmtTagName")
+
+            neat.prepare.instances.make_connection_on_exact_match(source, target, connection)
+            ```
         """
 
         subject_type, subject_predicate = self._get_type_and_property_uris(*source)
@@ -133,9 +261,67 @@ class InstancePrepareAPI:
         transformer = RelationshipAsEdgeTransformer(min_relationship_types, limit_per_type)
         self._state.instances.store.transform(transformer)
 
+    def convert_data_type(self, source: tuple[str, str], *, convert: Callable[[Any], Any] | None = None) -> None:
+        """Convert the data type of the given property.
+
+        This is, for example, useful when you have a boolean property that you want to convert to an enum.
+
+        Args:
+            source: The source of the conversion. A tuple of (type, property)
+                    where property is the property that should be converted.
+            convert: The function to use for the conversion. The function should take the value of the property
+                    as input and return the converted value. Default to assume you have a string that should be
+                    converted to int, float, bool, or datetime.
+
+        Example:
+            Convert a boolean property to a string:
+            ```python
+            neat.prepare.instances.convert_data_type(
+                ("TimeSeries", "isString"),
+                convert=lambda is_string: "string" if is_string else "numeric"
+            )
+            ```
+
+        """
+        subject_type, subject_predicate = self._get_type_and_property_uris(*source)
+
+        transformer = ConvertLiteral(subject_type, subject_predicate, convert)
+        self._state.instances.store.transform(transformer)
+
+    def property_to_type(self, source: tuple[str | None, str], type: str, new_property: str | None = None) -> None:
+        """Convert a property to a new type.
+
+        Args:
+            source: The source of the conversion. A tuple of (type, property)
+                    where property is the property that should be converted.
+                    You can pass (None, property) to covert all properties with the given name.
+            type: The new type of the property.
+            new_property: Add the identifier as a new property. If None, the new entity will not have a property.
+
+        Example:
+            Convert the property 'source' to SourceSystem
+            ```python
+            neat.prepare.instances.property_to_type(
+                (None, "source"), "SourceSystem"
+            )
+            ```
+        """
+        subject_type: URIRef | None = None
+        if source[0] is not None:
+            subject_type, subject_predicate = self._get_type_and_property_uris(*source)  # type: ignore[arg-type, assignment]
+        else:
+            subject_predicate = self._state.instances.store.queries.property_uri(source[1])[0]
+
+        transformer = LiteralToEntity(subject_type, subject_predicate, type, new_property)
+        self._state.instances.store.transform(transformer)
+
 
 @session_class_wrapper
 class DataModelPrepareAPI:
+    """Operations to perform on a data model as part of a workflow before writing the data model
+    to a desired destination.
+    """
+
     def __init__(self, client: NeatClient | None, state: SessionState, verbose: bool) -> None:
         self._client = client
         self._state = state
@@ -205,13 +391,15 @@ class DataModelPrepareAPI:
             move_connections: If True, the connections will be moved to the new data model.
 
         !!! note "Enterprise Data Model Creation"
+
             Always create an enterprise data model from a Cognite Data Model as this will
             assure all the Cognite Data Fusion applications to run smoothly, such as
-            - Search
-            - Atlas AI
-            - ...
+                - Search
+                - Atlas AI
+                - ...
 
         !!! note "Move Connections"
+
             If you want to move the connections to the new data model, set the move_connections
             to True. This will move the connections to the new data model and use new model
             views as the source and target views.
@@ -262,6 +450,7 @@ class DataModelPrepareAPI:
             dummy_property: The dummy property to use as placeholder for the views in the new data model.
 
         !!! note "Solution Data Model Mode"
+
             The read-only solution model will only be able to read from the existing containers
             from the enterprise data model, therefore the solution data model will not have
             containers in the solution data model space. Meaning the solution data model views

cognite/neat/_session/_read.py

@@ -1,7 +1,7 @@
 import tempfile
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Literal, cast
 
 from cognite.client.data_classes.data_modeling import DataModelId, DataModelIdentifier
 
@@ -11,7 +11,7 @@ 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._rules import importers
+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
@@ -27,6 +27,8 @@ from .exceptions import NeatSessionError, session_class_wrapper
 
 @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:
         self._state = state
         self._verbose = verbose
@@ -68,6 +70,11 @@ class BaseReadAPI:
 
 @session_class_wrapper
 class CDFReadAPI(BaseReadAPI):
+    """Reads from CDF Data Models.
+    Use the `.data_model()` method to load a CDF Data Model to the knowledge graph.
+
+    """
+
     def __init__(self, state: SessionState, client: NeatClient | None, verbose: bool) -> None:
         super().__init__(state, client, verbose)
         self.classic = CDFClassicAPI(state, client, verbose)
@@ -79,6 +86,18 @@ class CDFReadAPI(BaseReadAPI):
         return self._client
 
     def data_model(self, data_model_id: DataModelIdentifier) -> IssueList:
+        """Reads a Data Model from CDF to the knowledge graph.
+
+        Args:
+            data_model_id: Tuple of strings with the id of a CDF Data Model.
+            Notation as follows (<name_of_space>, <name_of_data_model>, <data_model_version>)
+
+        Example:
+            ```python
+            neat.read.cdf.data_model(("example_data_model_space", "EXAMPLE_DATA_MODEL", "v1"))
+            ```
+        """
+
         data_model_id = DataModelId.load(data_model_id)
 
         if not data_model_id.version:
@@ -113,6 +132,11 @@ class CDFReadAPI(BaseReadAPI):
 
 @session_class_wrapper
 class CDFClassicAPI(BaseReadAPI):
+    """Reads from the Classic Data Model from CDF.
+    Use the `.graph()` method to load CDF core resources to the knowledge graph.
+
+    """
+
     @property
     def _get_client(self) -> NeatClient:
         if self._client is None:
@@ -124,12 +148,12 @@ class CDFClassicAPI(BaseReadAPI):
 
         The Classic Graph consists of the following core resource type.
 
-        Classic Node CDF Resources:
-         - Assets
-         - TimeSeries
-         - Sequences
-         - Events
-         - Files
+        !!! note "Classic Node CDF Resources"
+             - Assets
+             - TimeSeries
+             - Sequences
+             - Events
+             - Files
 
         All the classic node CDF resources can have one or more connections to one or more assets. This
         will match a direct relationship in the data modeling of CDF.
@@ -144,13 +168,12 @@ class CDFClassicAPI(BaseReadAPI):
         This extractor will extract the classic CDF graph into Neat starting from either a data set or a root asset.
 
         It works as follows:
-
-        1. Extract all core nodes (assets, time series, sequences, events, files) filtered by the given data set or
-           root asset.
-        2. Extract all relationships starting from any of the extracted core nodes.
-        3. Extract all core nodes that are targets of the relationships that are not already extracted.
-        4. Extract all labels that are connected to the extracted core nodes/relationships.
-        5. Extract all data sets that are connected to the extracted core nodes/relationships.
+            1. Extract all core nodes (assets, time series, sequences, events, files) filtered by the given data set or
+               root asset.
+            2. Extract all relationships starting from any of the extracted core nodes.
+            3. Extract all core nodes that are targets of the relationships that are not already extracted.
+            4. Extract all labels that are connected to the extracted core nodes/relationships.
+            5. Extract all data sets that are connected to the extracted core nodes/relationships.
 
         Args:
             root_asset_external_id: The external id of the root asset
@@ -168,7 +191,29 @@ class CDFClassicAPI(BaseReadAPI):
 
 @session_class_wrapper
 class ExcelReadAPI(BaseReadAPI):
+    """Reads a Neat Excel Rules sheet to the graph store. The rules sheet may stem from an Information architect,
+    or a DMS Architect.
+
+    Args:
+        io: file path to the Excel sheet
+
+    Example:
+        ```python
+        neat.read.excel("information_or_dms_rules_sheet.xlsx")
+        ```
+    """
+
+    def __init__(self, state: SessionState, client: NeatClient | None, verbose: bool) -> None:
+        super().__init__(state, client, verbose)
+        self.examples = ExcelExampleAPI(state, client, 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,
+        or a DMS Architect.
+
+        Args:
+            io: file path to the Excel sheet
+        """
         reader = NeatReader.create(io)
         start = datetime.now(timezone.utc)
         if not isinstance(reader, PathReader):
@@ -190,8 +235,45 @@ class ExcelReadAPI(BaseReadAPI):
         return input_rules.issues
 
 
+@session_class_wrapper
+class ExcelExampleAPI(BaseReadAPI):
+    """Used as example for reading some data model into the NeatSession."""
+
+    @property
+    def pump_example(self) -> IssueList:
+        """Reads the Nordic 44 knowledge graph into the NeatSession graph store."""
+        start = datetime.now(timezone.utc)
+        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
+
+
 @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".
+
+    Example:
+        ```python
+        neat.read.yaml("path_to_toolkit_yamls")
+        ```
+    """
+
     def __call__(self, io: Any, format: Literal["neat", "toolkit"] = "neat") -> IssueList:
         reader = NeatReader.create(io)
         if not isinstance(reader, PathReader):
@@ -242,6 +324,23 @@ class YamlReadAPI(BaseReadAPI):
 
 @session_class_wrapper
 class CSVReadAPI(BaseReadAPI):
+    """Reads a csv that contains a column to use as primary key which will be the unique identifier for the type of
+    data you want to read in. Ex. a csv can hold information about assets, and their identifiers are specified in
+    a "ASSET_TAG" column.
+
+    Args:
+        io: file path or url to the csv
+        type: string that specifies what type of data the csv contains. For instance "Asset" or "Equipment"
+        primary_key: string name of the column that should be used as the unique identifier for each row of data
+
+    Example:
+        ```python
+        type_described_in_table = "Turbine"
+        column_with_identifier = "UNIQUE_TAG_NAME"
+        neat.read.csv("url_or_path_to_csv_file", type=type_described_in_table, primary_key=column_with_identifier)
+        ```
+    """
+
     def __call__(self, io: Any, type: str, primary_key: str) -> None:
         reader = NeatReader.create(io)
         if isinstance(reader, HttpFileReader):
@@ -263,6 +362,13 @@ class CSVReadAPI(BaseReadAPI):
 
 @session_class_wrapper
 class XMLReadAPI(BaseReadAPI):
+    """Reads an XML file that is either of DEXPI or AML format.
+
+    Args:
+        io: file path or url to the XML
+        format: can be either "dexpi" or "aml" are the currenly supported XML source types.
+    """
+
     def __call__(
         self,
         io: Any,
@@ -289,6 +395,16 @@ class XMLReadAPI(BaseReadAPI):
             raise NeatValueError("Only support XML files of DEXPI format at the moment.")
 
     def dexpi(self, path):
+        """Reads a DEXPI file into the NeatSession.
+
+        Args:
+            io: file path or url to the DEXPI file
+
+        Example:
+            ```python
+            neat.read.xml.dexpi("url_or_path_to_dexpi_file")
+            ```
+        """
         engine = import_engine()
         engine.set.format = "dexpi"
         engine.set.file = path
@@ -296,6 +412,16 @@ class XMLReadAPI(BaseReadAPI):
         self._state.instances.store.write(extractor)
 
     def aml(self, path):
+        """Reads an AML file into NeatSession.
+
+        Args:
+            io: file path or url to the AML file
+
+        Example:
+            ```python
+            neat.read.xml.aml("url_or_path_to_aml_file")
+            ```
+        """
         engine = import_engine()
         engine.set.format = "aml"
         engine.set.file = path
@@ -305,11 +431,27 @@ class XMLReadAPI(BaseReadAPI):
 
 @session_class_wrapper
 class RDFReadAPI(BaseReadAPI):
+    """Reads an RDF source into NeatSession. Supported sources are "ontology" or "imf".
+
+    Args:
+        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)
         self.examples = RDFExamples(state)
 
     def ontology(self, io: Any) -> IssueList:
+        """Reads an OWL ontology source into NeatSession.
+
+        Args:
+            io: file path or url to the OWL file
+
+        Example:
+            ```python
+            neat.read.rdf.ontology("url_or_path_to_owl_source")
+            ```
+        """
         start = datetime.now(timezone.utc)
         reader = NeatReader.create(io)
         if not isinstance(reader, PathReader):
@@ -331,6 +473,16 @@ class RDFReadAPI(BaseReadAPI):
         return input_rules.issues
 
     def imf(self, io: Any) -> IssueList:
+        """Reads IMF Types provided as SHACL shapes into NeatSession.
+
+        Args:
+            io: file path or url to the IMF file
+
+        Example:
+            ```python
+            neat.read.rdf.imf("url_or_path_to_imf_source")
+            ```
+        """
         start = datetime.now(timezone.utc)
         reader = NeatReader.create(io)
         if not isinstance(reader, PathReader):
@@ -360,16 +512,20 @@ class RDFReadAPI(BaseReadAPI):
         if type is None:
             type = object_wizard()
 
-        if type.lower() == "Data Model".lower():
+        type = type.lower()
+
+        if type == "data model":
             source = source or rdf_dm_wizard("What type of data model is the RDF?")
-            if source == "Ontology":
+            source = cast(str, source).lower()  # type: ignore
+
+            if source == "ontology":
                 return self.ontology(io)
-            elif source == "IMF":
+            elif source == "imf types":
                 return self.imf(io)
             else:
-                raise ValueError(f"Expected ontology, imf or instances, got {source}")
+                raise ValueError(f"Expected ontology, imf types or instances, got {source}")
 
-        elif type.lower() == "Instances".lower():
+        elif type == "instances":
             reader = NeatReader.create(io)
             if not isinstance(reader, PathReader):
                 raise NeatValueError("Only file paths are supported for RDF files")
@@ -380,11 +536,15 @@ class RDFReadAPI(BaseReadAPI):
             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."""
+
     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))
         return IssueList()

cognite/neat/_session/_set.py

@@ -12,12 +12,22 @@ from .exceptions import NeatSessionError, session_class_wrapper
 
 @session_class_wrapper
 class SetAPI:
+    """Used to change the name of the data model from a data model id defined by neat to a user specified name."""
+
     def __init__(self, state: SessionState, verbose: bool) -> None:
         self._state = state
         self._verbose = verbose
 
     def data_model_id(self, new_model_id: dm.DataModelId | tuple[str, str, str]) -> None:
-        """Sets the data model ID of the latest verified data model."""
+        """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>).
+
+        Example:
+            Set a new data model id:
+            ```python
+            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