cognite-neat 0.108.0__py3-none-any.whl → 0.109.0__py3-none-any.whl

cognite/neat/_session/_state.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass, field
+from pathlib import Path
 from typing import Literal, cast
 
 from cognite.neat._client import NeatClient
@@ -6,58 +6,46 @@ from cognite.neat._graph.extractors import KnowledgeGraphExtractor
 from cognite.neat._issues import IssueList
 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._rules.transformers import (
+    VerifiedRulesTransformer,
+)
 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
+from .exceptions import NeatSessionError, _session_method_wrapper
 
 
 class SessionState:
-    def __init__(self, store_type: Literal["memory", "oxigraph"], client: NeatClient | None = None) -> None:
-        self.instances = InstancesState(store_type)
+    def __init__(
+        self,
+        store_type: Literal["memory", "oxigraph"],
+        storage_path: Path | None = None,
+        client: NeatClient | None = None,
+    ) -> None:
+        self.instances = InstancesState(store_type, storage_path=storage_path)
         self.rule_store = NeatRulesStore()
         self.last_reference: DMSRules | InformationRules | None = None
         self.client = client
         self.quoted_source_identifiers = False
 
-    def rule_transform(self, *transformer: RulesTransformer) -> IssueList:
+    def rule_transform(self, *transformer: VerifiedRulesTransformer) -> IssueList:
         if not transformer:
             raise NeatSessionError("No transformers provided.")
-        first_transformer = transformer[0]
 
-        # This should not be allowed to be done automatically
-        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
+        start = 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}"
+        last_entity = self.rule_store.provenance[-1].target_entity
+        issues.action = f"{start} → {last_entity.display_name}"
         issues.hint = "Use the .inspect.issues() for more details."
+        self.instances.store.add_rules(last_entity.information)
         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
+        issues = self.rule_store.import_rules(importer, client=self.client)
+        if self.rule_store.empty:
+            result = "failed"
+        else:
+            result = self.rule_store.provenance[-1].target_entity.display_name
         if isinstance(importer, InferenceImporter):
             issues.action = f"Inferred {result}"
         else:
@@ -74,30 +62,39 @@ class SessionState:
         return issues
 
 
-@dataclass
 class InstancesState:
-    store_type: Literal["memory", "oxigraph"]
-    issue_lists: list[IssueList] = field(default_factory=list)
-    outcome: list[UploadResultList] = field(default_factory=list)
-    _store: NeatGraphStore | None = field(init=False, default=None)
+    def __init__(
+        self,
+        store_type: Literal["memory", "oxigraph"],
+        storage_path: Path | None = None,
+    ) -> None:
+        self.store_type = store_type
+        self.storage_path = storage_path
+        self.issue_lists = IssueList()
+        self.outcome = UploadResultList()
 
-    @property
-    def store(self) -> NeatGraphStore:
-        if not self.has_store:
-            if self.store_type == "oxigraph":
-                self._store = NeatGraphStore.from_oxi_local_store()
-            else:
-                self._store = NeatGraphStore.from_memory_store()
-        return cast(NeatGraphStore, self._store)
+        # Ensure that error handling is done in the constructor
+        self.store = _session_method_wrapper(self._create_store, "NeatSession")()
+
+        if self.storage_path:
+            print("Remember to close neat session .close() once you are done to avoid oxigraph lock.")
+
+    def _create_store(self) -> NeatGraphStore:
+        if self.store_type == "oxigraph":
+            if self.storage_path:
+                self.storage_path.mkdir(parents=True, exist_ok=True)
+            return NeatGraphStore.from_oxi_local_store(storage_dir=self.storage_path)
+        else:
+            return NeatGraphStore.from_memory_store()
 
     @property
-    def has_store(self) -> bool:
-        return self._store is not None
+    def empty(self) -> bool:
+        return self.store.empty
 
     @property
     def last_outcome(self) -> UploadResultList:
         if not self.outcome:
             raise NeatSessionError(
-                "No outcome available. Try using [bold].to.cdf.instances[/bold] to upload a data minstances."
+                "No outcome available. Try using [bold].to.cdf.instances[/bold] to upload a data instance."
             )
-        return self.outcome[-1]
+        return cast(UploadResultList, self.outcome[-1])

cognite/neat/_session/_to.py

@@ -4,7 +4,7 @@ from collections.abc import Collection
 from pathlib import Path
 from typing import Any, Literal, overload
 
-from cognite.client.data_classes.data_modeling import SpaceApply
+from cognite.client import data_modeling as dm
 
 from cognite.neat._constants import COGNITE_MODELS
 from cognite.neat._graph import loaders
@@ -194,7 +194,10 @@ class CDFToAPI:
         self._state = state
         self._verbose = verbose
 
-    def instances(self, space: str | None = None) -> UploadResultList:
+    def instances(
+        self,
+        space: str | None = None,
+    ) -> UploadResultList:
         """Export the verified DMS instances to CDF.
 
         Args:
@@ -213,7 +216,7 @@ class CDFToAPI:
             raise NeatSessionError("Please provide a valid space name. {PATTERNS.space_compliance.pattern}")
 
         if not client.data_modeling.spaces.retrieve(space):
-            client.data_modeling.spaces.apply(SpaceApply(space=space))
+            client.data_modeling.spaces.apply(dm.SpaceApply(space=space))
 
         loader = loaders.DMSLoader.from_rules(
             self._state.rule_store.last_verified_dms_rules,
@@ -224,6 +227,7 @@ class CDFToAPI:
             # urllib.parse.unquote() on the load.
             unquote_external_ids=self._state.quoted_source_identifiers,
         )
+
         result = loader.load_into_cdf(client)
         self._state.instances.outcome.append(result)
         print("You can inspect the details with the .inspect.outcome.instances(...) method.")

cognite/neat/_session/exceptions.py

@@ -3,6 +3,7 @@ from collections.abc import Callable
 from typing import Any
 
 from cognite.neat._issues.errors import CDFMissingClientError, NeatImportError
+from cognite.neat._issues.errors._external import OxigraphStorageLockedError
 from cognite.neat._issues.errors._general import NeatValueError
 
 from ._collector import _COLLECTOR
@@ -34,7 +35,12 @@ def _session_method_wrapper(func: Callable, cls_name: str):
         except NeatSessionError as e:
             action = _get_action()
             print(f"{_PREFIX} Cannot {action}: {e}")
-        except (CDFMissingClientError, NeatImportError, NeatValueError) as e:
+        except (
+            CDFMissingClientError,
+            NeatImportError,
+            NeatValueError,
+            OxigraphStorageLockedError,
+        ) as e:
             print(f"{_PREFIX} {escape(e.as_message())}")
         except ModuleNotFoundError as e:
             if e.name == "neatengine":

cognite/neat/_store/_graph_store.py

@@ -17,6 +17,7 @@ from cognite.neat._graph.extractors import RdfFileExtractor, TripleExtractors
 from cognite.neat._graph.queries import Queries
 from cognite.neat._graph.transformers import Transformers
 from cognite.neat._issues import IssueList, catch_issues
+from cognite.neat._issues.errors import OxigraphStorageLockedError
 from cognite.neat._rules.analysis import InformationAnalysis
 from cognite.neat._rules.models import InformationRules
 from cognite.neat._rules.models.entities import ClassEntity
@@ -24,7 +25,7 @@ from cognite.neat._shared import InstanceType, Triple
 from cognite.neat._utils.auxiliary import local_import
 from cognite.neat._utils.rdf_ import add_triples_in_batch, remove_namespace_from_uri
 
-from ._provenance import Change, Provenance
+from ._provenance import Change, Entity, Provenance
 
 if sys.version_info < (3, 11):
     from typing_extensions import Self
@@ -59,7 +60,7 @@ class NeatGraphStore:
 
         _start = datetime.now(timezone.utc)
         self.dataset = dataset
-        self.provenance = Provenance(
+        self.provenance = Provenance[Entity](
             [
                 Change.record(
                     activity=f"{type(self).__name__}.__init__",
@@ -207,17 +208,12 @@ class NeatGraphStore:
         import oxrdflib
         import pyoxigraph
 
-        # Adding support for both oxigraph in-memory and file-based storage
-        for i in range(4):
-            try:
-                oxi_store = pyoxigraph.Store(path=str(storage_dir) if storage_dir else None)
-                break
-            except OSError as e:
-                if "lock" in str(e) and i < 3:
-                    continue
-                raise e
-        else:
-            raise Exception("Error initializing Oxigraph store")
+        try:
+            oxi_store = pyoxigraph.Store(path=str(storage_dir) if storage_dir else None)
+        except OSError as e:
+            if "lock" in str(e):
+                raise OxigraphStorageLockedError(filepath=cast(Path, storage_dir)) from e
+            raise e
 
         return cls(
             dataset=Dataset(
@@ -685,3 +681,8 @@ class NeatGraphStore:
     @property
     def named_graphs(self) -> list[URIRef]:
         return [cast(URIRef, context.identifier) for context in self.dataset.contexts()]
+
+    @property
+    def empty(self) -> bool:
+        """Cheap way to check if the graph store is empty."""
+        return not self.queries.has_data()

cognite/neat/_store/_provenance.py

@@ -20,7 +20,7 @@ import uuid
 from collections.abc import Iterable, Sequence
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Optional
+from typing import Generic, TypeVar
 
 from rdflib import PROV, RDF, Literal, URIRef
 
@@ -49,9 +49,31 @@ UNKNOWN_AGENT = Agent(acted_on_behalf_of="UNKNOWN", id_=DEFAULT_NAMESPACE["unkno
 @dataclass(frozen=True)
 class Entity:
     was_attributed_to: Agent
-    issues: IssueList = field(default_factory=IssueList)
-    was_generated_by: Optional["Activity"] = field(default=None, repr=False)
-    id_: URIRef = DEFAULT_NAMESPACE["graph-store"]
+    issues: IssueList
+    was_generated_by: "Activity | None" = field(repr=False)
+    id_: URIRef
+
+    @classmethod
+    def create_with_defaults(
+        cls,
+        was_attributed_to: Agent,
+        issues: IssueList | None = None,
+        was_generated_by: "Activity | None" = None,
+        id_: URIRef = DEFAULT_NAMESPACE["graph-store"],
+    ) -> "Entity":
+        return cls(
+            was_attributed_to=was_attributed_to,
+            issues=issues or IssueList(),
+            was_generated_by=was_generated_by,
+            id_=id_,
+        )
+
+    @classmethod
+    def create_new_unknown_entity(cls) -> "Entity":
+        return cls.create_with_defaults(
+            was_attributed_to=UNKNOWN_AGENT,
+            id_=DEFAULT_NAMESPACE[f"unknown-entity/{uuid.uuid4()}"],
+        )
 
     def as_triples(self) -> list[Triple]:
         output: list[tuple[URIRef, URIRef, Literal | URIRef]] = [
@@ -70,16 +92,10 @@ class Entity:
 
         return output
 
-    @classmethod
-    def new_unknown_entity(cls) -> "Entity":
-        return cls(
-            was_attributed_to=UNKNOWN_AGENT,
-            id_=DEFAULT_NAMESPACE[f"unknown-entity/{uuid.uuid4()}"],
-        )
-
 
-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"])
+T_Entity = TypeVar("T_Entity", bound=Entity)
+INSTANCES_ENTITY = Entity.create_with_defaults(was_attributed_to=NEAT_AGENT, id_=CDF_NAMESPACE["instances"])
+EMPTY_ENTITY = Entity.create_with_defaults(was_attributed_to=NEAT_AGENT, id_=DEFAULT_NAMESPACE["empty-entity"])
 
 
 @dataclass(frozen=True)
@@ -111,12 +127,12 @@ class Activity:
 
 
 @dataclass(frozen=True)
-class Change(FrozenNeatObject):
+class Change(FrozenNeatObject, Generic[T_Entity]):
     agent: Agent
     activity: Activity
-    target_entity: Entity
+    target_entity: T_Entity
     description: str
-    source_entity: Entity = field(default_factory=Entity.new_unknown_entity)
+    source_entity: Entity = field(default_factory=Entity.create_new_unknown_entity)
 
     def as_triples(self) -> list[Triple]:
         return (
@@ -127,7 +143,7 @@ class Change(FrozenNeatObject):
         )
 
     @classmethod
-    def record(cls, activity: str, start: datetime, end: datetime, description: str) -> "Change":
+    def record(cls, activity: str, start: datetime, end: datetime, description: str) -> "Change[Entity]":
         """User friendly method to record a change that occurred in the graph store."""
         agent = Agent()
         activity = Activity(
@@ -136,8 +152,8 @@ class Change(FrozenNeatObject):
             started_at_time=start,
             ended_at_time=end,
         )
-        target_entity = Entity(was_generated_by=activity, was_attributed_to=agent)
-        return cls(
+        target_entity = Entity.create_with_defaults(was_generated_by=activity, was_attributed_to=agent)
+        return Change(
             agent=agent,
             activity=activity,
             target_entity=target_entity,
@@ -154,7 +170,7 @@ class Change(FrozenNeatObject):
         }
 
 
-class Provenance(NeatList[Change]):
+class Provenance(NeatList[Change[T_Entity]]):
     def __init__(self, changes: Sequence[Change] | None = None):
         super().__init__(changes or [])