move cdm subset method under read

cognitedata · Mar 5, 2025 · baf0778 · baf0778
1 parent 526769f
commit baf0778
Show file tree

Hide file tree

Showing 3 changed files with 99 additions and 179 deletions.
diff --git a/cognite/neat/_session/_read.py b/cognite/neat/_session/_read.py
@@ -28,7 +28,12 @@
 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.models.entities._single_value import ViewEntity
 from cognite.neat._rules.transformers import ClassicPrepareCore
+from cognite.neat._rules.transformers._converters import (
+    ToEnterpriseModel,
+    _SubsetEditableCDMRules,
+)
 from cognite.neat._utils.reader import NeatReader
 
 from ._state import SessionState
@@ -114,6 +119,86 @@ def data_model(self, data_model_id: DataModelIdentifier) -> IssueList:
         importer = importers.DMSImporter.from_data_model_id(cast(NeatClient, self._state.client), data_model_id)
         return self._state.rule_import(importer)
 
+    def core_data_model(self, concepts: str | list[str]) -> IssueList:
+        """Subset the data model to the desired concepts.
+
+        Args:
+            concepts: The concepts to subset the data model to
+
+        Returns:
+            IssueList: A list of issues that occurred during the transformation.
+
+        Example:
+            Read the CogniteCore data model and reduce the data model to only the 'CogniteAsset' concept.
+            ```python
+            neat = NeatSession(CogniteClient())
+
+            neat.subset.data_model.core_data_model(concepts=["CogniteAsset", "CogniteEquipment"])
+            ```
+
+        !!! note "Bundle of actions"
+            This method is a helper method that bundles the following actions:
+            - Imports the latest version of Cognite's Core Data Model (CDM)
+            - Makes editable copy of the CDM concepts
+            - Subsets the copy to the desired concepts to desired set of concepts
+        """
+
+        concepts = concepts if isinstance(concepts, list | set) else [concepts]
+
+        self._state._raise_exception_if_condition_not_met(
+            "Subset Core Data Model",
+            empty_rules_store_required=True,
+            client_required=True,
+        )
+
+        warnings.filterwarnings("default")
+        ExperimentalFlags.core_data_model_subsetting.warn()
+
+        cdm_v1 = DataModelId.load(("cdf_cdm", "CogniteCore", "v1"))
+        importer: importers.DMSImporter = importers.DMSImporter.from_data_model_id(
+            cast(NeatClient, self._state.client), cdm_v1
+        )
+        issues = self._state.rule_import(importer)
+
+        if issues.has_errors:
+            return issues
+
+        cdm_rules = self._state.rule_store.last_verified_rules
+
+        issues.extend(
+            self._state.rule_transform(
+                ToEnterpriseModel(
+                    new_model_id=("my_space", "MyCDMSubset", "v1"),
+                    org_name="CopyOf",
+                    dummy_property="GUID",
+                    move_connections=True,
+                )
+            )
+        )
+
+        if issues.has_errors:
+            return issues
+
+        issues.extend(
+            self._state.rule_transform(
+                _SubsetEditableCDMRules(
+                    views={
+                        ViewEntity(
+                            space=cdm_v1.space,
+                            externalId=concept,
+                            version=cast(str, cdm_v1.version),
+                        )
+                        for concept in concepts
+                    }
+                )
+            )
+        )
+
+        if cdm_rules and not issues.has_errors:
+            self._state.last_reference = cdm_rules
+
+        return issues
+
     def graph(
         self,
         data_model_id: DataModelIdentifier,

diff --git a/cognite/neat/_session/_subset.py b/cognite/neat/_session/_subset.py
@@ -1,18 +1,9 @@
 import warnings
-from typing import cast
-
-from cognite.client.data_classes.data_modeling import DataModelId
 
 from cognite.neat._alpha import ExperimentalFlags
-from cognite.neat._client._api_client import NeatClient
 from cognite.neat._issues._base import IssueList
-from cognite.neat._rules import importers
 from cognite.neat._rules.models.entities._single_value import ClassEntity, ViewEntity
 from cognite.neat._rules.transformers import SubsetDMSRules, SubsetInformationRules
-from cognite.neat._rules.transformers._converters import (
-    ToEnterpriseModel,
-    _SubsetEditableCDMRules,
-)
 
 from ._state import SessionState
 from .exceptions import NeatSessionError, session_class_wrapper
@@ -30,17 +21,10 @@ class SubsetAPI:
 
     """
 
-    def __init__(self, state: SessionState):
-        self._state = state
-        self.data_model = DataModelSubsetAPI(state)
-
-
-@session_class_wrapper
-class DataModelSubsetAPI:
     def __init__(self, state: SessionState):
         self._state = state
 
-    def __call__(self, concepts: str | list[str]) -> IssueList:
+    def data_model(self, concepts: str | list[str]) -> IssueList:
         """Subset the data model to the desired concepts.
 
         Args:
@@ -96,83 +80,3 @@ def __call__(self, concepts: str | list[str]) -> IssueList:
             print(f"Subset to {after} concepts.")
 
         return issues
-
-    def core_data_model(self, concepts: str | list[str]) -> IssueList:
-        """Subset the data model to the desired concepts.
-
-        Args:
-            concepts: The concepts to subset the data model to.
-
-        Returns:
-            IssueList: A list of issues that occurred during the transformation.
-
-        Example:
-            Read the CogniteCore data model and reduce the data model to only the 'CogniteAsset' concept.
-            ```python
-            neat = NeatSession(CogniteClient())
-
-            neat.subset.data_model.core_data_model(concepts=["CogniteAsset", "CogniteEquipment"])
-            ```
-
-        !!! note "Bundle of actions"
-            This method is a helper method that bundles the following actions:
-            - Imports the latest version of Cognite's Core Data Model (CDM)
-            - Makes editable copy of the CDM concepts
-            - Subsets the copy to the desired concepts to desired set of concepts
-        """
-
-        concepts = concepts if isinstance(concepts, list | set) else [concepts]
-
-        self._state._raise_exception_if_condition_not_met(
-            "Subset Core Data Model",
-            empty_rules_store_required=True,
-            client_required=True,
-        )
-
-        warnings.filterwarnings("default")
-        ExperimentalFlags.core_data_model_subsetting.warn()
-
-        cdm_v1 = DataModelId.load(("cdf_cdm", "CogniteCore", "v1"))
-        importer: importers.DMSImporter = importers.DMSImporter.from_data_model_id(
-            cast(NeatClient, self._state.client), cdm_v1
-        )
-        issues = self._state.rule_import(importer)
-
-        if issues.has_errors:
-            return issues
-
-        cdm_rules = self._state.rule_store.last_verified_rules
-
-        issues.extend(
-            self._state.rule_transform(
-                ToEnterpriseModel(
-                    new_model_id=("my_space", "MyCDMSubset", "v1"),
-                    org_name="CopyOf",
-                    dummy_property="GUID",
-                    move_connections=True,
-                )
-            )
-        )
-
-        if issues.has_errors:
-            return issues
-
-        issues.extend(
-            self._state.rule_transform(
-                _SubsetEditableCDMRules(
-                    views={
-                        ViewEntity(
-                            space=cdm_v1.space,
-                            externalId=concept,
-                            version=cast(str, cdm_v1.version),
-                        )
-                        for concept in concepts
-                    }
-                )
-            )
-        )
-
-        if cdm_rules and not issues.has_errors:
-            self._state.last_reference = cdm_rules
-
-        return issues
diff --git a/docs/tutorials/data-modeling/core_extension_via_subset.ipynb b/docs/tutorials/data-modeling/core_extension_via_subset.ipynb
@@ -121,7 +121,10 @@
     "- `MetMast`\n",
     "\n",
     "\n",
-    "To simplify this process we have create a convenience method `neat.subset.data_model.core_data_model()` to create editable set of CDM concepts that we can extend.\n",
+    "To simplify this process we have create a convenience method `neat.read.cdf.core_data_model()` to create editable set of CDM concepts that we can extend.\n",
+    "\n",
+    "\n",
+    "\n",
     "Let's call this method and pass the list of desired CDM concepts:\n"
    ]
   },
@@ -182,13 +185,13 @@
        "</div><br />Hint: Use the .inspect.issues() for more details."
       ],
       "text/plain": [
-       "[NotNeatSupportedFilterWarning(view_id=ViewId(space='cdf_cdm', external_id='CognitePointCloudModel', version='v1')),\n",
-       " NotNeatSupportedFilterWarning(view_id=ViewId(space='cdf_cdm', external_id='Cognite360ImageStation', version='v1')),\n",
-       " NotNeatSupportedFilterWarning(view_id=ViewId(space='cdf_cdm', external_id='CogniteCADRevision', version='v1')),\n",
-       " NotNeatSupportedFilterWarning(view_id=ViewId(space='cdf_cdm', external_id='CognitePointCloudRevision', version='v1')),\n",
+       "[NotNeatSupportedFilterWarning(view_id=ViewId(space='cdf_cdm', external_id='Cognite360ImageStation', version='v1')),\n",
+       " NotNeatSupportedFilterWarning(view_id=ViewId(space='cdf_cdm', external_id='CognitePointCloudModel', version='v1')),\n",
        " NotNeatSupportedFilterWarning(view_id=ViewId(space='cdf_cdm', external_id='CogniteCADModel', version='v1')),\n",
+       " NotNeatSupportedFilterWarning(view_id=ViewId(space='cdf_cdm', external_id='CogniteCADRevision', version='v1')),\n",
+       " NotNeatSupportedFilterWarning(view_id=ViewId(space='cdf_cdm', external_id='Cognite360ImageCollection', version='v1')),\n",
        " NotNeatSupportedFilterWarning(view_id=ViewId(space='cdf_cdm', external_id='Cognite360ImageModel', version='v1')),\n",
-       " NotNeatSupportedFilterWarning(view_id=ViewId(space='cdf_cdm', external_id='Cognite360ImageCollection', version='v1'))]"
+       " NotNeatSupportedFilterWarning(view_id=ViewId(space='cdf_cdm', external_id='CognitePointCloudRevision', version='v1'))]"
       ]
      },
      "execution_count": 4,
@@ -197,7 +200,7 @@
     }
    ],
    "source": [
-    "neat.subset.data_model.core_data_model(\n",
+    "neat.read.cdf.core_data_model(\n",
     "    [\"CogniteAsset\", \"CogniteEquipment\", \"CogniteTimeSeries\", \"CogniteActivity\", \"CogniteDescribable\"]\n",
     ")"
    ]
@@ -210,7 +213,7 @@
     "> Do not get confused with potential warnings you get when reading CDM into `NeatSession`. The warnings just point to users that filters are used in CDM. We typically strongly advise against usage of filters as it is easy to make mistakes when setting them.\n",
     "\n",
     "\n",
-    "Let's now inspect content of NeatSession by calling `neat` which will give us a summary of the created data model in `NeatSessio`:\n"
+    "Let's now inspect content of NeatSession by calling `neat` which will give us a summary of the created data model in `NeatSession`:\n"
    ]
   },
   {
@@ -285,7 +288,7 @@
        "</div>"
       ],
       "text/plain": [
-       "<cognite.neat._session._base.NeatSession at 0x1259db8d0>"
+       "<cognite.neat._session._base.NeatSession at 0x12c3dbe90>"
       ]
      },
      "execution_count": 5,
@@ -407,7 +410,7 @@
        "</div>"
       ],
       "text/plain": [
-       "<cognite.neat._session._base.NeatSession at 0x1259db8d0>"
+       "<cognite.neat._session._base.NeatSession at 0x12c3dbe90>"
       ]
      },
      "execution_count": 7,
@@ -661,78 +664,6 @@
     "Finally let's push data model to CDF:"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "id": "8fbbde3e-caf6-436d-be44-9abe147d9197",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "You can inspect the details with the .inspect.outcome.data_model(...) method.\n"
-     ]
-    },
-    {
-     "data": {
-      "text/html": [
-       "<style type=\"text/css\">\n",
-       "</style>\n",
-       "<table id=\"T_6ba5a\">\n",
-       "  <thead>\n",
-       "    <tr>\n",
-       "      <th class=\"blank level0\" >&nbsp;</th>\n",
-       "      <th id=\"T_6ba5a_level0_col0\" class=\"col_heading level0 col0\" >name</th>\n",
-       "      <th id=\"T_6ba5a_level0_col1\" class=\"col_heading level0 col1\" >unchanged</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th id=\"T_6ba5a_level0_row0\" class=\"row_heading level0 row0\" >0</th>\n",
-       "      <td id=\"T_6ba5a_row0_col0\" class=\"data row0 col0\" >spaces</td>\n",
-       "      <td id=\"T_6ba5a_row0_col1\" class=\"data row0 col1\" >1</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th id=\"T_6ba5a_level0_row1\" class=\"row_heading level0 row1\" >1</th>\n",
-       "      <td id=\"T_6ba5a_row1_col0\" class=\"data row1 col0\" >containers</td>\n",
-       "      <td id=\"T_6ba5a_row1_col1\" class=\"data row1 col1\" >9</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th id=\"T_6ba5a_level0_row2\" class=\"row_heading level0 row2\" >2</th>\n",
-       "      <td id=\"T_6ba5a_row2_col0\" class=\"data row2 col0\" >views</td>\n",
-       "      <td id=\"T_6ba5a_row2_col1\" class=\"data row2 col1\" >9</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th id=\"T_6ba5a_level0_row3\" class=\"row_heading level0 row3\" >3</th>\n",
-       "      <td id=\"T_6ba5a_row3_col0\" class=\"data row3 col0\" >data_models</td>\n",
-       "      <td id=\"T_6ba5a_row3_col1\" class=\"data row3 col1\" >1</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th id=\"T_6ba5a_level0_row4\" class=\"row_heading level0 row4\" >4</th>\n",
-       "      <td id=\"T_6ba5a_row4_col0\" class=\"data row4 col0\" >nodes</td>\n",
-       "      <td id=\"T_6ba5a_row4_col1\" class=\"data row4 col1\" >0</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n"
-      ],
-      "text/plain": [
-       "[UploadResult(name='spaces', error_messages=[], issues=[], created=set(), upserted=set(), deleted=set(), changed=set(), unchanged={'wind_energy_space'}, skipped=set(), failed_created=set(), failed_upserted=set(), failed_changed=set(), failed_deleted=set(), failed_items=[]),\n",
-       " UploadResult(name='containers', error_messages=[], issues=[], created=set(), upserted=set(), deleted=set(), changed=set(), unchanged={ContainerId(space='wind_energy_space', external_id='Location'), ContainerId(space='wind_energy_space', external_id='WindTurbine'), ContainerId(space='wind_energy_space', external_id='CopyOfEquipment'), ContainerId(space='wind_energy_space', external_id='MetMast'), ContainerId(space='wind_energy_space', external_id='CopyOfTimeSeries'), ContainerId(space='wind_energy_space', external_id='CopyOfAsset'), ContainerId(space='wind_energy_space', external_id='WindFarm'), ContainerId(space='wind_energy_space', external_id='CopyOfActivity'), ContainerId(space='wind_energy_space', external_id='Substation')}, skipped=set(), failed_created=set(), failed_upserted=set(), failed_changed=set(), failed_deleted=set(), failed_items=[]),\n",
-       " UploadResult(name='views', error_messages=[], issues=[], created=set(), upserted=set(), deleted=set(), changed=set(), unchanged={ViewId(space='wind_energy_space', external_id='Location', version='v1'), ViewId(space='wind_energy_space', external_id='WindTurbine', version='v1'), ViewId(space='wind_energy_space', external_id='WindFarm', version='v1'), ViewId(space='wind_energy_space', external_id='CopyOfEquipment', version='v1'), ViewId(space='wind_energy_space', external_id='CopyOfAsset', version='v1'), ViewId(space='wind_energy_space', external_id='Substation', version='v1'), ViewId(space='wind_energy_space', external_id='CopyOfActivity', version='v1'), ViewId(space='wind_energy_space', external_id='MetMast', version='v1'), ViewId(space='wind_energy_space', external_id='CopyOfTimeSeries', version='v1')}, skipped=set(), failed_created=set(), failed_upserted=set(), failed_changed=set(), failed_deleted=set(), failed_items=[]),\n",
-       " UploadResult(name='data_models', error_messages=[], issues=[], created=set(), upserted=set(), deleted=set(), changed=set(), unchanged={DataModelId(space='wind_energy_space', external_id='TinyWindFarmModel', version='v1')}, skipped=set(), failed_created=set(), failed_upserted=set(), failed_changed=set(), failed_deleted=set(), failed_items=[]),\n",
-       " UploadResult(name='nodes', error_messages=[], issues=[], created=set(), upserted=set(), deleted=set(), changed=set(), unchanged=set(), skipped=set(), failed_created=set(), failed_upserted=set(), failed_changed=set(), failed_deleted=set(), failed_items=[])]"
-      ]
-     },
-     "execution_count": 10,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "neat.to.cdf.data_model()"
-   ]
-  },
   {
    "cell_type": "markdown",
    "id": "63fd6fba-b2cf-4c52-b988-891949a7ebbd",