cognite-neat 0.103.1__py3-none-any.whl → 0.105.0__py3-none-any.whl

cognite/neat/_session/_state.py

@@ -1,15 +1,14 @@
 from dataclasses import dataclass, field
 from typing import Literal, cast
 
-from rdflib import URIRef
-
 from cognite.neat._issues import IssueList
-from cognite.neat._rules._shared import JustRules, ReadRules, VerifiedRules
-from cognite.neat._rules.models.dms._rules import DMSRules
-from cognite.neat._rules.models.information._rules import InformationRules
-from cognite.neat._rules.models.information._rules_input import InformationInputRules
-from cognite.neat._store import NeatGraphStore
-from cognite.neat._store._provenance import Change, Provenance
+from cognite.neat._rules.importers import BaseImporter, InferenceImporter
+from cognite.neat._rules.models import DMSRules, InformationRules
+from cognite.neat._rules.transformers import RulesTransformer, ToExtensionModel
+from cognite.neat._store import NeatGraphStore, NeatRulesStore
+from cognite.neat._store._rules_store import ModelEntity
+from cognite.neat._utils.rdf_ import uri_display_name
+from cognite.neat._utils.text import humanize_collection
 from cognite.neat._utils.upload import UploadResultList
 
 from .exceptions import NeatSessionError
@@ -18,7 +17,48 @@ from .exceptions import NeatSessionError
 class SessionState:
     def __init__(self, store_type: Literal["memory", "oxigraph"]) -> None:
         self.instances = InstancesState(store_type)
-        self.data_model = DataModelState()
+        self.rule_store = NeatRulesStore()
+        self.last_reference: DMSRules | InformationRules | None = None
+
+    def rule_transform(self, *transformer: RulesTransformer) -> IssueList:
+        if not transformer:
+            raise NeatSessionError("No transformers provided.")
+        first_transformer = transformer[0]
+        pruned = self.rule_store.prune_until_compatible(first_transformer)
+        if pruned:
+            type_hint = first_transformer.transform_type_hint()
+            action = uri_display_name(first_transformer.agent.id_)
+            location = cast(ModelEntity, self.rule_store.provenance[-1].target_entity).display_name
+            expected = humanize_collection([hint.display_type_name() for hint in type_hint])  # type: ignore[attr-defined]
+            step_str = "step" if len(pruned) == 1 else "steps"
+            print(
+                f"The {action} actions expects a {expected}. "
+                f"Moving back {len(pruned)} {step_str} to the last {location}."
+            )
+        if (
+            any(isinstance(t, ToExtensionModel) for t in transformer)
+            and isinstance(self.rule_store.provenance[-1].target_entity, ModelEntity)
+            and isinstance(self.rule_store.provenance[-1].target_entity.result, DMSRules | InformationRules)
+        ):
+            self.last_reference = self.rule_store.provenance[-1].target_entity.result
+
+        start = cast(ModelEntity, self.rule_store.provenance[-1].target_entity).display_name
+        issues = self.rule_store.transform(*transformer)
+        end = cast(ModelEntity, self.rule_store.provenance[-1].target_entity).display_name
+        issues.action = f"{start} → {end}"
+        issues.hint = "Use the .inspect.issues() for more details."
+        return issues
+
+    def rule_import(self, importer: BaseImporter) -> IssueList:
+        issues = self.rule_store.import_(importer)
+        result = cast(ModelEntity, self.rule_store.provenance[-1].target_entity).display_name
+        if isinstance(importer, InferenceImporter):
+            issues.action = f"Inferred {result}"
+        else:
+            issues.action = f"Read {result}"
+        if issues:
+            issues.hint = "Use the .inspect.issues() for more details."
+        return issues
 
 
 @dataclass
@@ -48,101 +88,3 @@ class InstancesState:
                 "No outcome available. Try using [bold].to.cdf.instances[/bold] to upload a data minstances."
             )
         return self.outcome[-1]
-
-
-@dataclass
-class DataModelState:
-    _rules: dict[URIRef, ReadRules | JustRules | VerifiedRules] = field(init=False, default_factory=dict)
-    issue_lists: list[IssueList] = field(default_factory=list)
-    provenance: Provenance = field(default_factory=Provenance)
-    outcome: list[UploadResultList] = field(default_factory=list)
-
-    def write(self, rules: ReadRules | JustRules | VerifiedRules, change: Change) -> None:
-        if change.target_entity.id_ in self._rules:
-            raise NeatSessionError(f"Data model <{change.target_entity.id_}> already exists.")
-
-        else:
-            self._rules[change.target_entity.id_] = rules
-            self.provenance.append(change)
-
-    def read(self, id_: URIRef) -> ReadRules | JustRules | VerifiedRules:
-        if id_ not in self._rules:
-            raise NeatSessionError(f"Data model <{id_}> not found.")
-        return self._rules[id_]
-
-    @property
-    def unverified_rules(self) -> dict[URIRef, ReadRules]:
-        return {id_: rules for id_, rules in self._rules.items() if isinstance(rules, ReadRules)}
-
-    @property
-    def verified_rules(self) -> dict[URIRef, VerifiedRules]:
-        return {id_: rules for id_, rules in self._rules.items() if isinstance(rules, VerifiedRules)}
-
-    @property
-    def last_unverified_rule(self) -> tuple[URIRef, ReadRules]:
-        if not self.unverified_rules:
-            raise NeatSessionError("No data model available. Try using [bold].read[/bold] to load a data model.")
-        return next(reversed(self.unverified_rules.items()))
-
-    @property
-    def last_info_unverified_rule(self) -> tuple[URIRef, ReadRules]:
-        if self.unverified_rules:
-            for id_, rule in reversed(self.unverified_rules.items()):
-                if isinstance(rule.rules, InformationInputRules):
-                    return id_, rule
-
-        raise NeatSessionError("No data model available. Try using [bold].read[/bold] to load a data model.")
-
-    @property
-    def last_verified_rule(self) -> tuple[URIRef, VerifiedRules]:
-        if not self.verified_rules:
-            raise NeatSessionError(
-                "No data model available to verify. Try using  [bold].read[/bold] to load a data model."
-            )
-        return next(reversed(self.verified_rules.items()))
-
-    @property
-    def last_verified_information_rules(self) -> tuple[URIRef, InformationRules]:
-        if self.verified_rules:
-            for id_, rules in reversed(self.verified_rules.items()):
-                if isinstance(rules, InformationRules):
-                    return id_, rules
-
-        raise NeatSessionError(
-            "No verified information data model. Try using  [bold].verify()[/bold]"
-            " to convert unverified information model to verified information model."
-        )
-
-    @property
-    def last_verified_dms_rules(self) -> tuple[URIRef, DMSRules]:
-        if self.verified_rules:
-            for id_, rules in reversed(self.verified_rules.items()):
-                if isinstance(rules, DMSRules):
-                    return id_, rules
-
-        raise NeatSessionError(
-            'No verified DMS data model. Try using  [bold].convert("DMS")[/bold]'
-            " to convert verified information model to verified DMS model."
-        )
-
-    @property
-    def has_unverified_rules(self) -> bool:
-        return bool(self.unverified_rules)
-
-    @property
-    def has_verified_rules(self) -> bool:
-        return bool(self.verified_rules)
-
-    @property
-    def last_issues(self) -> IssueList:
-        if not self.issue_lists:
-            raise NeatSessionError("No issues available. Try using [bold].verify()[/bold] to verify a data model.")
-        return self.issue_lists[-1]
-
-    @property
-    def last_outcome(self) -> UploadResultList:
-        if not self.outcome:
-            raise NeatSessionError(
-                "No outcome available. Try using [bold].to.cdf.data_model[/bold] to upload a data model."
-            )
-        return self.outcome[-1]

cognite/neat/_session/_to.py

@@ -5,13 +5,14 @@ from typing import Any, Literal, overload
 from cognite.client.data_classes.data_modeling import SpaceApply
 
 from cognite.neat._client import NeatClient
+from cognite.neat._constants import COGNITE_MODELS
 from cognite.neat._graph import loaders
-from cognite.neat._issues import IssueList, catch_warnings
 from cognite.neat._rules import exporters
 from cognite.neat._rules._constants import PATTERNS
 from cognite.neat._rules._shared import VerifiedRules
 from cognite.neat._rules.exporters._rules2dms import Component
-from cognite.neat._utils.upload import UploadResultCore, UploadResultList
+from cognite.neat._rules.models.dms import DMSMetadata
+from cognite.neat._utils.upload import UploadResultList
 
 from ._state import SessionState
 from .exceptions import NeatSessionError, session_class_wrapper
@@ -32,45 +33,59 @@ class ToAPI:
     def excel(
         self,
         io: Any,
-        model: Literal["dms", "information", "logical", "physical"] | None,
+        include_reference: bool = True,
     ) -> None:
         """Export the verified data model to Excel.
 
         Args:
             io: The file path or file-like object to write the Excel file to.
-            model: The format of the data model to export. Defaults to None.
+            include_reference: If True, the reference data model will be included. Defaults to True.
+                Note that this only applies if you have created the data model using the
+                .to_enterprise(), .to_solution(), or .to_data_product() methods.
 
         Example:
             Export information model to excel rules sheet
             ```python
             information_rules_file_name = "information_rules.xlsx"
-            neat.to.excel(information_rules_file_name, model="information")
+            neat.to.excel(information_rules_file_name)
             ```
 
         Example:
-            Export data model to excel rules sheet
+            Read CogniteCore model, convert it to an enterprise model, and export it to an excel file
             ```python
+            client = CogniteClient()
+            neat = NeatSession(client)
+
+            neat.read.cdf(("cdf_cdm", "CogniteCore", "v1"))
+            neat.verify()
+            neat.prepare.data_model.to_enterprise(
+                data_model_id=("sp_doctrino_space", "ExtensionCore", "v1"),
+                org_name="MyOrg",
+                move_connections=True
+            )
             dms_rules_file_name = "dms_rules.xlsx"
-            neat.to.excel(information_rules_file_name, model="dms")
+            neat.to.excel(dms_rules_file_name, include_reference=True)
             ```
         """
-        exporter = exporters.ExcelExporter(styling="maximal")
-        rules: VerifiedRules
-        if model == "information" or model == "logical":
-            rules = self._state.data_model.last_verified_information_rules[1]
-        elif model == "dms" or model == "physical":
-            rules = self._state.data_model.last_verified_dms_rules[1]
-        else:
-            rules = self._state.data_model.last_verified_rule[1]
-
-        exporter.export_to_file(rules, Path(io))
-        return None
+        reference_rules_with_prefix: tuple[VerifiedRules, str] | None = None
+        if include_reference and self._state.last_reference:
+            if (
+                isinstance(self._state.last_reference.metadata, DMSMetadata)
+                and self._state.last_reference.metadata.as_data_model_id() in COGNITE_MODELS
+            ):
+                prefix = "CDM"
+            else:
+                prefix = "Ref"
+            reference_rules_with_prefix = self._state.last_reference, prefix
+
+        exporter = exporters.ExcelExporter(styling="maximal", reference_rules_with_prefix=reference_rules_with_prefix)
+        return self._state.rule_store.export_to_file(exporter, Path(io))
 
     @overload
-    def yaml(self, io: None, format: Literal["neat"] = "neat") -> str: ...
+    def yaml(self, io: None, format: Literal["neat"] = "neat", skip_system_spaces: bool = True) -> str: ...
 
     @overload
-    def yaml(self, io: Any, format: Literal["neat", "toolkit"] = "neat") -> None: ...
+    def yaml(self, io: Any, format: Literal["neat", "toolkit"] = "neat", skip_system_spaces: bool = True) -> None: ...
 
     def yaml(
         self, io: Any | None = None, format: Literal["neat", "toolkit"] = "neat", skip_system_spaces: bool = True
@@ -112,22 +127,22 @@ class ToAPI:
         """
         if format == "neat":
             exporter = exporters.YAMLExporter()
-            last_verified = self._state.data_model.last_verified_rule[1]
             if io is None:
-                return exporter.export(last_verified)
+                return self._state.rule_store.export(exporter)
 
-            exporter.export_to_file(last_verified, Path(io))
+            self._state.rule_store.export_to_file(exporter, Path(io))
         elif format == "toolkit":
             if io is None or not isinstance(io, str | Path):
                 raise NeatSessionError(
                     "Please provide a zip file or directory path to write the YAML files to."
                     "This is required for the 'toolkit' format."
                 )
-            dms_rule = self._state.data_model.last_verified_dms_rules[1]
             user_path = Path(io)
             if user_path.suffix == "" and not user_path.exists():
                 user_path.mkdir(parents=True)
-            exporters.DMSExporter(remove_cdf_spaces=skip_system_spaces).export_to_file(dms_rule, user_path)
+            self._state.rule_store.export_to_file(
+                exporters.DMSExporter(remove_cdf_spaces=skip_system_spaces), user_path
+            )
         else:
             raise NeatSessionError("Please provide a valid format. 'neat' or 'toolkit'")
 
@@ -154,9 +169,9 @@ class CDFToAPI:
         if not self._client:
             raise NeatSessionError("No CDF client provided!")
 
-        space = space or f"{self._state.data_model.last_verified_dms_rules[1].metadata.space}_instances"
+        space = space or f"{self._state.rule_store.last_verified_dms_rules.metadata.space}_instances"
 
-        if space and space == self._state.data_model.last_verified_dms_rules[1].metadata.space:
+        if space and space == self._state.rule_store.last_verified_dms_rules.metadata.space:
             raise NeatSessionError("Space for instances must be different from the data model space.")
         elif not PATTERNS.space_compliance.match(str(space)):
             raise NeatSessionError("Please provide a valid space name. {PATTERNS.space_compliance.pattern}")
@@ -165,7 +180,7 @@ class CDFToAPI:
             self._client.data_modeling.spaces.apply(SpaceApply(space=space))
 
         loader = loaders.DMSLoader.from_rules(
-            self._state.data_model.last_verified_dms_rules[1],
+            self._state.rule_store.last_verified_dms_rules,
             self._state.instances.store,
             instance_space=space,
             client=self._client,
@@ -207,10 +222,6 @@ class CDFToAPI:
         if not self._client:
             raise NeatSessionError("No client provided!")
 
-        conversion_issues = IssueList(action="to.cdf.data_model")
-        with catch_warnings(conversion_issues):
-            result = exporter.export_to_cdf(self._state.data_model.last_verified_dms_rules[1], self._client, dry_run)
-        result.insert(0, UploadResultCore(name="schema", issues=conversion_issues))
-        self._state.data_model.outcome.append(result)
+        result = self._state.rule_store.export_to_cdf(exporter, self._client, dry_run)
         print("You can inspect the details with the .inspect.outcome.data_model(...) method.")
         return result

cognite/neat/_shared.py

@@ -1,11 +1,17 @@
+import sys
 from abc import abstractmethod
-from collections.abc import Hashable, Sequence
+from collections.abc import Hashable, Iterator, Sequence
 from dataclasses import dataclass
-from typing import Any, TypeAlias, TypeVar
+from typing import Any, SupportsIndex, TypeAlias, TypeVar, overload
 
 import pandas as pd
 from rdflib import Literal, URIRef
 
+if sys.version_info <= (3, 11):
+    from typing_extensions import Self
+else:
+    from typing import Self
+
 T_ID = TypeVar("T_ID", bound=Hashable)
 
 
@@ -51,6 +57,21 @@ class NeatList(list, Sequence[T_NeatObject]):
     def _repr_html_(self) -> str:
         return self.to_pandas()._repr_html_()  # type: ignore[operator]
 
+    # Implemented to get correct type hints
+    def __iter__(self) -> Iterator[T_NeatObject]:
+        return super().__iter__()
+
+    @overload
+    def __getitem__(self, index: SupportsIndex) -> T_NeatObject: ...
+
+    @overload
+    def __getitem__(self, index: slice) -> Self: ...
+
+    def __getitem__(self, index: SupportsIndex | slice, /) -> T_NeatObject | Self:
+        if isinstance(index, slice):
+            return type(self)(super().__getitem__(index))
+        return super().__getitem__(index)
+
 
 Triple: TypeAlias = tuple[URIRef, URIRef, Literal | URIRef]
 InstanceType: TypeAlias = URIRef

cognite/neat/_store/_provenance.py

@@ -20,14 +20,12 @@ import uuid
 from collections.abc import Iterable, Sequence
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Generic, Optional, TypeVar
+from typing import Optional
 
-from cognite.client.data_classes.data_modeling import DataModelId, DataModelIdentifier
 from rdflib import PROV, RDF, Literal, URIRef
 
 from cognite.neat._constants import CDF_NAMESPACE, DEFAULT_NAMESPACE
 from cognite.neat._issues import IssueList
-from cognite.neat._rules._shared import JustRules, ReadRules, VerifiedRules
 from cognite.neat._shared import FrozenNeatObject, NeatList, Triple
 
 
@@ -72,45 +70,6 @@ class Entity:
 
         return output
 
-    @classmethod
-    def from_data_model_id(cls, data_model_id: DataModelIdentifier) -> "Entity":
-        data_model_id = DataModelId.load(data_model_id)
-
-        return cls(
-            was_attributed_to=CDF_AGENT,
-            id_=CDF_NAMESPACE[
-                f"dms/data-model/{data_model_id.space}/{data_model_id.external_id}/{data_model_id.version}"
-            ],
-        )
-
-    @classmethod
-    def from_rules(
-        cls,
-        rules: ReadRules | JustRules | VerifiedRules,
-        agent: Agent | None = None,
-        activity: "Activity | None" = None,
-    ) -> "Entity":
-        agent = agent or UNKNOWN_AGENT
-        if isinstance(rules, VerifiedRules):
-            return cls(
-                was_attributed_to=agent,
-                was_generated_by=activity,
-                id_=rules.metadata.identifier,
-            )
-
-        elif isinstance(rules, ReadRules | JustRules) and rules.rules is not None:
-            return cls(
-                was_attributed_to=agent,
-                was_generated_by=activity,
-                id_=rules.rules.metadata.identifier,
-            )
-        else:
-            return cls(
-                was_attributed_to=agent,
-                was_generated_by=activity,
-                id_=DEFAULT_NAMESPACE["unknown-entity"],
-            )
-
     @classmethod
     def new_unknown_entity(cls) -> "Entity":
         return cls(
@@ -119,15 +78,8 @@ class Entity:
         )
 
 
-T_Result = TypeVar("T_Result")
-
-
-@dataclass(frozen=True)
-class ModelEntity(Entity, Generic[T_Result]):
-    result: T_Result | None = None
-
-
 INSTANCES_ENTITY = Entity(was_attributed_to=NEAT_AGENT, id_=CDF_NAMESPACE["instances"])
+EMPTY_ENTITY = Entity(was_attributed_to=NEAT_AGENT, id_=DEFAULT_NAMESPACE["empty-entity"])
 
 
 @dataclass(frozen=True)
@@ -192,34 +144,6 @@ class Change(FrozenNeatObject):
             description=description,
         )
 
-    @classmethod
-    def from_rules_activity(
-        cls,
-        rules: ReadRules | JustRules | VerifiedRules,
-        agent: Agent,
-        start: datetime,
-        end: datetime,
-        description: str,
-        source_entity: Entity | None = None,
-    ) -> "Change":
-        source_entity = source_entity or Entity.new_unknown_entity()
-        activity = Activity(
-            started_at_time=start,
-            ended_at_time=end,
-            was_associated_with=agent,
-            used=source_entity,
-        )
-
-        target_entity = Entity.from_rules(rules, agent, activity)
-
-        return cls(
-            agent=agent,
-            activity=activity,
-            target_entity=target_entity,
-            description=description,
-            source_entity=source_entity,
-        )
-
     def dump(self, aggregate: bool = True) -> dict[str, str]:
         return {
             "Source Entity": self.source_entity.id_,
@@ -230,11 +154,8 @@ class Change(FrozenNeatObject):
         }
 
 
-T_Change = TypeVar("T_Change", bound=Change)
-
-
 class Provenance(NeatList[Change]):
-    def __init__(self, changes: Sequence[T_Change] | None = None):
+    def __init__(self, changes: Sequence[Change] | None = None):
         super().__init__(changes or [])
 
     def activity_took_place(self, activity: str) -> bool: