cognite-neat 0.94.0__py3-none-any.whl → 0.96.0__py3-none-any.whl

cognite/neat/_session/_base.py

@@ -1,19 +1,20 @@
 from typing import Literal, cast
 
-import pandas as pd
 from cognite.client import CogniteClient
 
+from cognite.neat import _version
 from cognite.neat._issues import IssueList
 from cognite.neat._rules import importers
 from cognite.neat._rules._shared import ReadRules
 from cognite.neat._rules.models import DMSRules
-from cognite.neat._rules.models._base_input import InputComponent
 from cognite.neat._rules.models.information._rules import InformationRules
 from cognite.neat._rules.models.information._rules_input import InformationInputRules
 from cognite.neat._rules.transformers import ConvertToRules, VerifyAnyRules
 
+from ._inspect import InspectAPI
 from ._prepare import PrepareAPI
 from ._read import ReadAPI
+from ._set import SetAPI
 from ._show import ShowAPI
 from ._state import SessionState
 from ._to import ToAPI
@@ -35,6 +36,12 @@ class NeatSession:
         self.to = ToAPI(self._state, client, verbose)
         self.prepare = PrepareAPI(self._state, verbose)
         self.show = ShowAPI(self._state)
+        self.set = SetAPI(self._state, verbose)
+        self.inspect = InspectAPI(self._state)
+
+    @property
+    def version(self) -> str:
+        return _version.__version__
 
     def verify(self) -> IssueList:
         output = VerifyAnyRules("continue").try_transform(self._state.input_rule)
@@ -42,6 +49,9 @@ class NeatSession:
             self._state.verified_rules.append(output.rules)
             if isinstance(output.rules, InformationRules):
                 self._state.store.add_rules(output.rules)
+        self._state.issue_lists.append(output.issues)
+        if output.issues:
+            print("You can inspect the issues with the .inspect attribute.")
         return output.issues
 
     def convert(self, target: Literal["dms"]) -> None:
@@ -72,13 +82,10 @@ class NeatSession:
 
         output = []
         if state.input_rules and not state.verified_rules:
-            metadata = cast(InputComponent, state.input_rule.rules.metadata)  # type: ignore[union-attr]
-            table = pd.DataFrame([metadata.dump()]).T._repr_html_()  # type: ignore[operator]
-            output.append(f"<H2>Raw Data Model</H2><br />{table}")
+            output.append(f"<H2>Unverified Data Model</H2><br />{state.input_rule.rules._repr_html_()}")  # type: ignore
 
         if state.verified_rules:
-            table = pd.DataFrame([state.last_verified_rule.metadata.model_dump()]).T._repr_html_()  # type: ignore[operator]
-            output.append(f"<H2>Data Model</H2><br />{table}")
+            output.append(f"<H2>Verified Data Model</H2><br />{state.last_verified_rule._repr_html_()}")  # type: ignore
 
         if state.has_store:
             output.append(f"<H2>Instances</H2> {state.store._repr_html_()}")

cognite/neat/_session/_inspect.py

@@ -0,0 +1,89 @@
+import difflib
+from typing import Literal, overload
+
+import pandas as pd
+
+from cognite.neat._constants import IN_NOTEBOOK
+from cognite.neat._issues import IssueList
+
+from ._state import SessionState
+from .exceptions import intercept_session_exceptions
+
+
+@intercept_session_exceptions
+class InspectAPI:
+    def __init__(self, state: SessionState) -> None:
+        self._state = state
+        self.issues = InspectIssues(state)
+
+    @property
+    def properties(self) -> pd.DataFrame:
+        """Returns the properties of the current data model."""
+        return self._state.last_verified_rule.properties.to_pandas()
+
+
+@intercept_session_exceptions
+class InspectIssues:
+    """Inspect issues of the current data model."""
+
+    def __init__(self, state: SessionState) -> None:
+        self._state = state
+
+    @overload
+    def __call__(
+        self,
+        search: str | None = None,
+        return_dataframe: Literal[True] = (False if IN_NOTEBOOK else True),  # type: ignore[assignment]
+    ) -> pd.DataFrame: ...
+
+    @overload
+    def __call__(
+        self,
+        search: str | None = None,
+        return_dataframe: Literal[False] = (False if IN_NOTEBOOK else True),  # type: ignore[assignment]
+    ) -> None: ...
+
+    def __call__(
+        self,
+        search: str | None = None,
+        return_dataframe: bool = (False if IN_NOTEBOOK else True),  # type: ignore[assignment]
+    ) -> pd.DataFrame | None:
+        """Returns the issues of the current data model."""
+        issues = self._state.last_issues
+        if not issues:
+            self._print("No issues found.")
+
+        if issues and search is not None:
+            unique_types = {type(issue).__name__ for issue in issues}
+            closest_match = set(difflib.get_close_matches(search, unique_types))
+            issues = IssueList([issue for issue in issues if type(issue).__name__ in closest_match])
+
+        if IN_NOTEBOOK:
+            from IPython.display import Markdown, display
+
+            issue_str = "\n".join(
+                [f"  * **{type(issue).__name__}**: {issue.as_message(include_type=False)}" for issue in issues]
+            )
+            message = f"### {len(issues)} issues found\n\n{issue_str}"
+            display(Markdown(message))
+
+        if return_dataframe:
+            return issues.to_pandas()
+        return None
+
+    def _print(self, message: str) -> None:
+        if IN_NOTEBOOK:
+            from IPython.display import Markdown, display
+
+            display(Markdown(message))
+        else:
+            print(message)
+
+    def __repr__(self) -> str:
+        return self.__repr_html__()
+
+    def __repr_html__(self) -> str:
+        return (
+            "Inspect issues by calling .inspect.issues() or "
+            "search for specific issues by calling .inspect.issues('MyTypeWarning')."
+        )

cognite/neat/_session/_prepare.py

@@ -9,8 +9,10 @@ from cognite.neat._rules.models.information._rules_input import InformationInput
 from cognite.neat._rules.transformers import ReduceCogniteModel, ToCompliantEntities, ToExtension
 
 from ._state import SessionState
+from .exceptions import intercept_session_exceptions
 
 
+@intercept_session_exceptions
 class PrepareAPI:
     def __init__(self, state: SessionState, verbose: bool) -> None:
         self._state = state
@@ -18,6 +20,7 @@ class PrepareAPI:
         self.data_model = DataModelPrepareAPI(state, verbose)
 
 
+@intercept_session_exceptions
 class DataModelPrepareAPI:
     def __init__(self, state: SessionState, verbose: bool) -> None:
         self._state = state
@@ -35,25 +38,79 @@ class DataModelPrepareAPI:
                 )
             )
 
-    def to_extension(self, new_data_model_id: DataModelIdentifier, org_name: str | None = None) -> None:
-        """Uses the current data model as a basis to extend from.
+    def to_enterprise(
+        self,
+        data_model_id: DataModelIdentifier,
+        org_name: str = "My",
+        dummy_property: str = "GUID",
+    ) -> None:
+        """Uses the current data model as a basis to create enterprise data model
 
         Args:
-            new_data_model_id: The new data model that is extending the current data model.
-            org_name: Organization name to use for the views in the new data model. This is required if you are
-                creating an extension from a Cognite Data Model.
+            data_model_id: The enterprise data model id that is being created
+            org_name: Organization name to use for the views in the enterprise data model.
+            dummy_property: The dummy property to use as placeholder for the views in 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
+            - ...
+
+        """
+        if dms := self._state.last_verified_dms_rules:
+            output = ToExtension(
+                new_model_id=data_model_id,
+                org_name=org_name,
+                type_="enterprise",
+                dummy_property=dummy_property,
+            ).transform(dms)
+            self._state.verified_rules.append(output.rules)
+
+    def to_solution(
+        self,
+        data_model_id: DataModelIdentifier,
+        org_name: str = "My",
+        mode: Literal["read", "write"] = "read",
+        dummy_property: str = "dummy",
+    ) -> None:
+        """Uses the current data model as a basis to create solution data model
+
+        Args:
+            data_model_id: The solution data model id that is being created.
+            org_name: Organization name to use for the views in the new data model.
+            mode: The mode of the solution data model. Can be either "read" or "write".
+            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
+            will be read-only.
+
+            The write mode will have additional containers in the solution data model space,
+            allowing in addition to reading through the solution model views, also writing to
+            the containers in the solution data model space.
 
         """
         if dms := self._state.last_verified_dms_rules:
-            output = ToExtension(new_data_model_id, org_name).transform(dms)
+            output = ToExtension(
+                new_model_id=data_model_id,
+                org_name=org_name,
+                type_="solution",
+                mode=mode,
+                dummy_property=dummy_property,
+            ).transform(dms)
             self._state.verified_rules.append(output.rules)
 
-    def reduce(self, drop: Collection[Literal["3D", "Annotation", "BaseViews"]]) -> None:
+    def reduce(self, drop: Collection[Literal["3D", "Annotation", "BaseViews"] | str]) -> None:
         """This is a special method that allow you to drop parts of the data model.
         This only applies to Cognite Data Models.
 
         Args:
-            drop: Which parts of the data model to drop.
+            drop: What to drop from the data model. The values 3D, Annotation, and BaseViews are special values that
+                drops multiple views at once. You can also pass externalIds of views to drop individual views.
 
         """
         if dms := self._state.last_verified_dms_rules:

cognite/neat/_session/_read.py

@@ -7,13 +7,16 @@ from cognite.client.data_classes.data_modeling import DataModelIdentifier
 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._shared import ReadRules
 
 from ._state import SessionState
 from ._wizard import NeatObjectType, RDFFileType, object_wizard, rdf_dm_wizard
+from .exceptions import intercept_session_exceptions
 
 
+@intercept_session_exceptions
 class ReadAPI:
     def __init__(self, state: SessionState, client: CogniteClient | None, verbose: bool) -> None:
         self._state = state
@@ -23,6 +26,7 @@ class ReadAPI:
         self.excel = ExcelReadAPI(state, client, verbose)
 
 
+@intercept_session_exceptions
 class BaseReadAPI:
     def __init__(self, state: SessionState, client: CogniteClient | None, verbose: bool) -> None:
         self._state = state
@@ -44,20 +48,44 @@ class BaseReadAPI:
         elif isinstance(io, Path):
             return io
         else:
-            raise ValueError(f"Expected str or Path, got {type(io)}")
+            raise NeatValueError(f"Expected str or Path, got {type(io)}")
 
 
+@intercept_session_exceptions
 class CDFReadAPI(BaseReadAPI):
-    def data_model(self, data_model_id: DataModelIdentifier) -> IssueList:
+    def __init__(self, state: SessionState, client: CogniteClient | None, verbose: bool) -> None:
+        super().__init__(state, client, verbose)
+        self.classic = CDFClassicAPI(state, client, verbose)
+
+    @property
+    def _get_client(self) -> CogniteClient:
         if self._client is None:
-            raise ValueError("No client provided. Please provide a client to read a data model.")
+            raise NeatValueError("No client provided. Please provide a client to read a data model.")
+        return self._client
 
-        importer = importers.DMSImporter.from_data_model_id(self._client, data_model_id)
+    def data_model(self, data_model_id: DataModelIdentifier) -> IssueList:
+        importer = importers.DMSImporter.from_data_model_id(self._get_client, data_model_id)
         input_rules = importer.to_rules()
         self._store_rules(data_model_id, input_rules, "CDF")
         return input_rules.issues
 
 
+@intercept_session_exceptions
+class CDFClassicAPI(BaseReadAPI):
+    @property
+    def _get_client(self) -> CogniteClient:
+        if self._client is None:
+            raise ValueError("No client provided. Please provide a client to read a data model.")
+        return self._client
+
+    def assets(self, root_asset_external_id: str) -> None:
+        extractor = extractors.AssetsExtractor.from_hierarchy(self._get_client, root_asset_external_id)
+        self._state.store.write(extractor)
+        if self._verbose:
+            print(f"Asset hierarchy {root_asset_external_id} read successfully")
+
+
+@intercept_session_exceptions
 class ExcelReadAPI(BaseReadAPI):
     def __call__(self, io: Any) -> IssueList:
         filepath = self._return_filepath(io)
@@ -66,6 +94,7 @@ class ExcelReadAPI(BaseReadAPI):
         return input_rules.issues
 
 
+@intercept_session_exceptions
 class RDFReadAPI(BaseReadAPI):
     def __init__(self, state: SessionState, client: CogniteClient | None, verbose: bool) -> None:
         super().__init__(state, client, verbose)

cognite/neat/_session/_set.py

@@ -0,0 +1,23 @@
+from cognite.client import data_modeling as dm
+
+from cognite.neat._rules.transformers import SetIDDMSModel
+
+from ._state import SessionState
+from .exceptions import intercept_session_exceptions
+
+
+@intercept_session_exceptions
+class SetAPI:
+    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."""
+        if dms := self._state.last_verified_dms_rules:
+            output = SetIDDMSModel(new_model_id).transform(dms)
+            self._state.verified_rules.append(output.rules)
+            if self._verbose:
+                print(f"Data model ID set to {new_model_id}")
+        else:
+            print("No verified DMS data model available")