feat: added option for keeping reference

cognitedata · Jan 18, 2025 · 519adcf · 519adcf
1 parent cf49dbd
commit 519adcf
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 4 deletions.
diff --git a/cognite/neat/_rules/transformers/_converters.py b/cognite/neat/_rules/transformers/_converters.py
@@ -41,7 +41,7 @@
 )
 from cognite.neat._rules.models._rdfpath import Entity as RDFPathEntity
 from cognite.neat._rules.models._rdfpath import RDFPath, SingleProperty
-from cognite.neat._rules.models.data_types import AnyURI, DataType, String
+from cognite.neat._rules.models.data_types import AnyURI, DataType, File, String, Timeseries
 from cognite.neat._rules.models.dms import DMSMetadata, DMSProperty, DMSValidation, DMSView
 from cognite.neat._rules.models.dms._rules import DMSContainer
 from cognite.neat._rules.models.entities import (
@@ -823,11 +823,21 @@ class ClassicPrepareCore(RulesTransformer[InformationRules, InformationRules]):
     - ClassicTimeseries.isString from boolean to string
     - Add class ClassicSourceSystem, and update all source properties from string to ClassicSourceSystem.
     - Rename externalId properties to classicExternalId
-    - Renames the Relationship.sourceExternaId and Relationship.targetExternalId to startNode and endNode
+    - Renames the Relationship.sourceExternalId and Relationship.targetExternalId to startNode and endNode
+    - If reference_timeseries is True, the classicExternalId property of the TimeSeries class will change type
+      from string to timeseries.
+    - If reference_files is True, the classicExternalId property of the File class will change type from string to file.
     """
 
-    def __init__(self, instance_namespace: Namespace) -> None:
+    def __init__(
+        self,
+        instance_namespace: Namespace,
+        reference_timeseries: bool = False,
+        reference_files: bool = False,
+    ) -> None:
         self.instance_namespace = instance_namespace
+        self.reference_timeseries = reference_timeseries
+        self.reference_files = reference_files
 
     @property
     def description(self) -> str:
@@ -851,6 +861,10 @@ def transform(self, rules: InformationRules) -> InformationRules:
                 prop.value_type = ClassEntity(prefix=prefix, suffix="ClassicSourceSystem")
             elif prop.property_ == "externalId":
                 prop.property_ = "classicExternalId"
+                if self.reference_timeseries and prop.class_.suffix == "ClassicTimeSeries":
+                    prop.value_type = Timeseries()
+                elif self.reference_files and prop.class_.suffix == "ClassicFile":
+                    prop.value_type = File()
             elif prop.property_ == "sourceExternalId" and prop.class_.suffix == "ClassicRelationship":
                 prop.property_ = "startNode"
             elif prop.property_ == "targetExternalId" and prop.class_.suffix == "ClassicRelationship":

diff --git a/cognite/neat/_session/_read.py b/cognite/neat/_session/_read.py
@@ -181,6 +181,18 @@ def graph(
             neat.read.cdf.graph("root_asset_external_id")
             ```
         """
+        return self._graph(
+            root_asset_external_id, limit_per_type, identifier, reference_timeseries=False, reference_files=False
+        )
+
+    def _graph(
+        self,
+        root_asset_external_id: str,
+        limit_per_type: int | None = None,
+        identifier: Literal["id", "externalId"] = "id",
+        reference_timeseries: bool = False,
+        reference_files: bool = False,
+    ) -> IssueList:
         namespace = CLASSIC_CDF_NAMESPACE
         extractor = extractors.ClassicGraphExtractor(
             self._get_client,
@@ -208,7 +220,9 @@ def graph(
             LiteralToEntity(None, namespace["source"], "ClassicSourceSystem", "name"),
         )
         # Updating the information model.
-        prepare_issues = self._state.rule_store.transform(ClassicPrepareCore(namespace))
+        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