cognite-neat 0.78.4__py3-none-any.whl → 0.79.0__py3-none-any.whl

cognite/neat/graph/loaders/_rdf2dms.py

@@ -0,0 +1,280 @@
+import itertools
+import json
+from collections import defaultdict
+from collections.abc import Iterable, Sequence
+from pathlib import Path
+from typing import Any
+
+import yaml
+from cognite.client import CogniteClient
+from cognite.client import data_modeling as dm
+from cognite.client.data_classes.capabilities import Capability, DataModelInstancesAcl
+from cognite.client.data_classes.data_modeling import ViewId
+from cognite.client.data_classes.data_modeling.views import SingleEdgeConnection
+from cognite.client.exceptions import CogniteAPIError
+from pydantic import ValidationInfo, create_model, field_validator
+from pydantic.main import Model
+
+from cognite.neat.graph._tracking import LogTracker, Tracker
+from cognite.neat.graph.issues import loader as loader_issues
+from cognite.neat.graph.stores import NeatGraphStoreBase
+from cognite.neat.issues import NeatIssue, NeatIssueList
+from cognite.neat.rules.models import DMSRules
+from cognite.neat.rules.models.data_types import _DATA_TYPE_BY_DMS_TYPE
+from cognite.neat.utils.upload import UploadDiffsID
+
+from ._base import CDFLoader
+
+
+class DMSLoader(CDFLoader[dm.InstanceApply]):
+    def __init__(
+        self,
+        graph_store: NeatGraphStoreBase,
+        data_model: dm.DataModel[dm.View] | None,
+        instance_space: str,
+        class_by_view_id: dict[ViewId, str] | None = None,
+        creat_issues: Sequence[NeatIssue] | None = None,
+        tracker: type[Tracker] | None = None,
+    ):
+        super().__init__(graph_store)
+        self.data_model = data_model
+        self.instance_space = instance_space
+        self.class_by_view_id = class_by_view_id or {}
+        self._issues = NeatIssueList[NeatIssue](creat_issues or [])
+        self._tracker: type[Tracker] = tracker or LogTracker
+
+    @classmethod
+    def from_data_model_id(
+        cls,
+        client: CogniteClient,
+        data_model_id: dm.DataModelId,
+        graph_store: NeatGraphStoreBase,
+        instance_space: str,
+    ) -> "DMSLoader":
+        issues: list[NeatIssue] = []
+        data_model: dm.DataModel[dm.View] | None = None
+        try:
+            data_model = client.data_modeling.data_models.retrieve(data_model_id, inline_views=True).latest_version()
+        except Exception as e:
+            issues.append(loader_issues.MissingDataModelError(identifier=repr(data_model_id), reason=str(e)))
+
+        return cls(graph_store, data_model, instance_space, {}, issues)
+
+    @classmethod
+    def from_rules(cls, rules: DMSRules, graph_store: NeatGraphStoreBase, instance_space: str) -> "DMSLoader":
+        issues: list[NeatIssue] = []
+        data_model: dm.DataModel[dm.View] | None = None
+        try:
+            data_model = rules.as_schema().as_read_model()
+        except Exception as e:
+            issues.append(
+                loader_issues.FailedConvertError(
+                    identifier=rules.metadata.as_identifier(), target_format="read DMS model", reason=str(e)
+                )
+            )
+        return cls(graph_store, data_model, instance_space, {}, issues)
+
+    def _load(self, stop_on_exception: bool = False) -> Iterable[dm.InstanceApply | NeatIssue]:
+        if self._issues.has_errors and stop_on_exception:
+            raise self._issues.as_exception()
+        elif self._issues.has_errors:
+            yield from self._issues
+            return
+        if not self.data_model:
+            # There should already be an error in this case.
+            return
+        view_ids = [repr(v.as_id()) for v in self.data_model.views]
+        tracker = self._tracker(type(self).__name__, view_ids, "views")
+        for view in self.data_model.views:
+            view_id = view.as_id()
+            tracker.start(repr(view_id))
+            pydantic_cls, edge_by_properties, issues = self._create_validation_classes(view)  # type: ignore[var-annotated]
+            yield from issues
+            tracker.issue(issues)
+            class_name = self.class_by_view_id.get(view.as_id(), view.external_id)
+            triples = self.graph_store.queries.triples_of_type_instances(class_name)
+            for identifier, properties in _triples2dictionary(triples).items():
+                try:
+                    yield self._create_node(identifier, properties, pydantic_cls, view_id)
+                except ValueError as e:
+                    error = loader_issues.InvalidInstanceError(type_="node", identifier=identifier, reason=str(e))
+                    tracker.issue(error)
+                    if stop_on_exception:
+                        raise error.as_exception() from e
+                    yield error
+                yield from self._create_edges(identifier, properties, edge_by_properties, tracker)
+            tracker.finish(repr(view_id))
+
+    def write_to_file(self, filepath: Path) -> None:
+        if filepath.suffix not in [".json", ".yaml", ".yml"]:
+            raise ValueError(f"File format {filepath.suffix} is not supported")
+        dumped: dict[str, list] = {"nodes": [], "edges": [], "issues": []}
+        for item in self.load(stop_on_exception=False):
+            key = {
+                dm.NodeApply: "nodes",
+                dm.EdgeApply: "edges",
+                NeatIssue: "issues",
+            }.get(type(item))
+            if key is None:
+                # This should never happen, and is a bug in neat
+                raise ValueError(f"Item {item} is not supported. This is a bug in neat please report it.")
+            dumped[key].append(item.dump())
+        with filepath.open("w", encoding=self._encoding, newline=self._new_line) as f:
+            if filepath.suffix == ".json":
+                json.dump(dumped, f, indent=2)
+            else:
+                yaml.safe_dump(dumped, f, sort_keys=False)
+
+    def _create_validation_classes(
+        self, view: dm.View
+    ) -> tuple[type[Model], dict[str, dm.EdgeConnection], NeatIssueList]:
+        issues = NeatIssueList[NeatIssue]()
+        field_definitions: dict[str, tuple[type, Any]] = {}
+        edge_by_property: dict[str, dm.EdgeConnection] = {}
+        direct_relation_by_property: dict[str, dm.DirectRelation] = {}
+        for prop_name, prop in view.properties.items():
+            if isinstance(prop, dm.EdgeConnection):
+                edge_by_property[prop_name] = prop
+            if isinstance(prop, dm.MappedProperty):
+                if isinstance(prop.type, dm.DirectRelation):
+                    direct_relation_by_property[prop_name] = prop.type
+                    python_type: Any = dict
+                else:
+                    data_type = _DATA_TYPE_BY_DMS_TYPE.get(prop.type._type)
+                    if not data_type:
+                        issues.append(
+                            loader_issues.InvalidClassWarning(
+                                class_name=repr(view.as_id()),
+                                reason=f"Unknown data type for property {prop_name}: {prop.type._type}",
+                            )
+                        )
+                        continue
+                    python_type = data_type.python
+                if prop.type.is_list:
+                    python_type = list[python_type]
+                default_value: Any = prop.default_value
+                if prop.nullable:
+                    python_type = python_type | None
+                else:
+                    default_value = ...
+
+                field_definitions[prop_name] = (python_type, default_value)
+
+        def parse_list(cls, value: Any, info: ValidationInfo) -> list[str]:
+            if isinstance(value, list) and cls.model_fields[info.field_name].annotation is not list:
+                if len(value) == 1:
+                    return value[0]
+                raise ValueError(f"Got multiple values for {info.field_name}: {value}")
+            return value
+
+        validators: dict[str, classmethod] = {"parse_list": field_validator("*", mode="before")(parse_list)}  # type: ignore[dict-item,arg-type]
+        if direct_relation_by_property:
+
+            def parse_direct_relation(cls, value: list, info: ValidationInfo) -> dict | list[dict]:
+                # We validate above that we only get one value for single direct relations.
+                if cls.model_fields[info.field_name].annotation is list:
+                    return [{"space": self.instance_space, "externalId": v} for v in value]
+                elif value:
+                    return {"space": self.instance_space, "externalId": value[0]}
+                return {}
+
+            validators["parse_direct_relation"] = field_validator(*direct_relation_by_property.keys(), mode="before")(  # type: ignore[assignment]
+                parse_direct_relation  # type: ignore[arg-type]
+            )
+
+        pydantic_cls = create_model(view.external_id, __validators__=validators, **field_definitions)  # type: ignore[arg-type, call-overload]
+        return pydantic_cls, edge_by_property, issues
+
+    def _create_node(
+        self, identifier: str, properties: dict[str, list[str]], pydantic_cls: type[Model], view_id: dm.ViewId
+    ) -> dm.InstanceApply:
+        created = pydantic_cls.model_validate(properties)
+
+        return dm.NodeApply(
+            space=self.instance_space,
+            external_id=identifier,
+            # type=#RDF type
+            sources=[dm.NodeOrEdgeData(source=view_id, properties=dict(created.model_dump().items()))],
+        )
+
+    def _create_edges(
+        self,
+        identifier: str,
+        properties: dict[str, list[str]],
+        edge_by_properties: dict[str, dm.EdgeConnection],
+        tracker: Tracker,
+    ) -> Iterable[dm.EdgeApply | NeatIssue]:
+        for prop, values in properties.items():
+            if prop not in edge_by_properties:
+                continue
+            edge = edge_by_properties[prop]
+            if isinstance(edge, SingleEdgeConnection) and len(values) > 1:
+                error = loader_issues.InvalidInstanceError(
+                    type_="edge",
+                    identifier=identifier,
+                    reason=f"Multiple values for single edge {edge}. Expected only one.",
+                )
+                tracker.issue(error)
+                yield error
+            for target in values:
+                yield dm.EdgeApply(
+                    space=self.instance_space,
+                    external_id=f"{identifier}.{prop}.{target}",
+                    type=edge.type,
+                    start_node=dm.DirectRelationReference(self.instance_space, identifier),
+                    end_node=dm.DirectRelationReference(self.instance_space, target),
+                )
+
+    def _get_required_capabilities(self) -> list[Capability]:
+        return [
+            DataModelInstancesAcl(
+                actions=[
+                    DataModelInstancesAcl.Action.Write,
+                    DataModelInstancesAcl.Action.Write_Properties,
+                    DataModelInstancesAcl.Action.Read,
+                ],
+                scope=DataModelInstancesAcl.Scope.SpaceID([self.instance_space]),
+            )
+        ]
+
+    def _upload_to_cdf(
+        self,
+        client: CogniteClient,
+        items: list[dm.InstanceApply],
+        return_diffs: bool,
+        dry_run: bool,
+        read_issues: NeatIssueList,
+    ) -> UploadDiffsID:
+        result = UploadDiffsID(name=type(self).__name__, issues=read_issues)
+        try:
+            nodes = [item for item in items if isinstance(item, dm.NodeApply)]
+            edges = [item for item in items if isinstance(item, dm.EdgeApply)]
+            upserted = client.data_modeling.instances.apply(
+                nodes,
+                edges,
+                auto_create_end_nodes=True,
+                auto_create_start_nodes=True,
+                skip_on_version_conflict=True,
+            )
+        except CogniteAPIError as e:
+            result.error_messages.append(str(e))
+            result.failed.append([repr(instance.as_id()) for instance in items])  # type: ignore[arg-type, attr-defined]
+        else:
+            for instance in itertools.chain(upserted.nodes, upserted.edges):
+                if instance.was_modified and instance.created_time == instance.last_updated_time:
+                    result.created.append(repr(instance.as_id()))
+                elif instance.was_modified:
+                    result.changed.append(repr(instance.as_id()))
+                else:
+                    result.unchanged.append(repr(instance.as_id()))
+        return result if return_diffs else result.as_upload_result_ids()  # type: ignore[return-value]
+
+
+def _triples2dictionary(
+    triples: Iterable[tuple[str, str, str]],
+) -> dict[str, dict[str, list[str]]]:
+    """Converts list of triples to dictionary"""
+    values_by_property_by_identifier: dict[str, dict[str, list[str]]] = defaultdict(lambda: defaultdict(list))
+    for id_, property_, value in triples:
+        values_by_property_by_identifier[id_][property_].append(value)
+    return values_by_property_by_identifier

cognite/neat/graph/stores/_base.py

@@ -3,17 +3,22 @@ import sys
 import time
 from abc import ABC, abstractmethod
 from collections.abc import Iterable, Iterator
+from datetime import datetime
 from pathlib import Path
 from typing import Literal, TypeAlias, cast
 
 import pandas as pd
+import pytz
 from prometheus_client import Gauge, Summary
-from rdflib import Graph, Namespace, URIRef
+from rdflib import RDF, Graph, Namespace, URIRef
 from rdflib.query import Result, ResultRow
 
 from cognite.neat.constants import DEFAULT_NAMESPACE, PREFIXES
 from cognite.neat.graph.models import Triple
 from cognite.neat.graph.stores._rdf_to_graph import rdf_file_to_graph
+from cognite.neat.utils import remove_namespace
+
+from ._provenance import Change, Provenance
 
 if sys.version_info >= (3, 11):
     pass
@@ -48,6 +53,7 @@ class NeatGraphStoreBase(ABC):
         namespace: Namespace = DEFAULT_NAMESPACE,
         prefixes: dict = PREFIXES,
     ):
+        _start = datetime.now(pytz.utc)
         self.graph = graph or Graph()
         self.base_prefix: str = base_prefix
         self.namespace: Namespace = namespace
@@ -62,6 +68,16 @@ class NeatGraphStoreBase(ABC):
         self.internal_storage_dir_orig: Path | None = None
         self.storage_dirs_to_delete: list[Path] = []
         self.queries = _Queries(self)
+        self.provenance = Provenance(
+            [
+                Change.record(
+                    activity=f"{type(self).__name__}.__init__",
+                    start=_start,
+                    end=datetime.now(pytz.utc),
+                    description="Initialize graph store",
+                )
+            ]
+        )
 
     @abstractmethod
     def _set_graph(self) -> None:
@@ -109,9 +125,9 @@ class NeatGraphStoreBase(ABC):
 
         if base_prefix:
             self.base_prefix = base_prefix
-
-        self.graph.bind(self.base_prefix, self.namespace)
-        logging.info("Adding prefix %s with namespace %s", self.base_prefix, self.namespace)
+        if self.base_prefix:
+            self.graph.bind(self.base_prefix, self.namespace)
+            logging.info("Adding prefix %s with namespace %s", self.base_prefix, self.namespace)
         logging.info("Graph initialized")
 
     def reinitialize_graph(self):
@@ -361,3 +377,17 @@ class _Queries:
         logging.info(query)
         # Select queries gives an iterable of result rows
         return cast(list[ResultRow], list(self.store.query(query)))
+
+    def triples_of_type_instances(self, rdf_type: str) -> list[tuple[str, str, str]]:
+        """Get all triples of a given type.
+
+        This method assumes the graph has been transformed into the default namespace.
+        """
+        query = (
+            f"SELECT ?instance ?prop ?value "
+            f"WHERE {{ ?instance a <{self.store.namespace[rdf_type]}> . ?instance ?prop ?value . }} order by ?instance"
+        )
+        result = self.store.query(query)
+
+        # We cannot include the RDF.type in case there is a neat:type property
+        return [remove_namespace(*triple) for triple in result if triple[1] != RDF.type]  # type: ignore[misc, index]

cognite/neat/graph/stores/_provenance.py

@@ -0,0 +1,99 @@
+# we will use prov-o to represent the provenance of the neat graph store
+# basically tracking changes that occur in the graph store
+# prov-o use concepts of Agent, Activity and Entity to represent provenance
+# where in case of neat we have:
+# Agent: triples extractors, graph enhancers, contextualizers, etc.
+# Activity: write triple, add connection, etc.
+# Entity: neat graph store
+
+
+import uuid
+from collections import UserList
+from collections.abc import Sequence
+from dataclasses import dataclass
+from datetime import datetime
+from typing import TypeVar
+
+from rdflib import PROV, RDF, Literal, URIRef
+
+from cognite.neat.constants import DEFAULT_NAMESPACE
+
+
+@dataclass(frozen=True)
+class Agent:
+    id_: URIRef = DEFAULT_NAMESPACE.agent
+    acted_on_behalf_of: str = "NEAT"
+
+    def as_triples(self):
+        return [
+            (self.id_, RDF.type, PROV[type(self).__name__]),
+            (self.id_, PROV.actedOnBehalfOf, self.acted_on_behalf_of),
+        ]
+
+
+@dataclass(frozen=True)
+class Activity:
+    was_associated_with: Agent
+    ended_at_time: datetime
+    started_at_time: datetime
+    used: str  # this would be set to for example Extractor, Enhancer, Contextualizer, etc.
+    id_: URIRef = DEFAULT_NAMESPACE[f"activity-{uuid.uuid4()}"]
+
+    def as_triples(self):
+        return [
+            (self.id_, RDF.type, PROV[type(self).__name__]),
+            (self.id_, PROV.wasAssociatedWith, self.was_associated_with.id_),
+            (self.id_, PROV.startedAtTime, Literal(self.started_at_time)),
+            (self.id_, PROV.endedAtTime, Literal(self.ended_at_time)),
+            (self.id_, PROV.used, self.used),
+        ]
+
+
+@dataclass(frozen=True)
+class Entity:
+    was_generated_by: Activity
+    was_attributed_to: Agent
+    id_: URIRef = DEFAULT_NAMESPACE["graph-store"]
+
+    def as_triples(self):
+        return [
+            (self.id_, RDF.type, PROV[type(self).__name__]),
+            (self.id_, PROV.wasGeneratedBy, self.was_generated_by.id_),
+            (self.id_, PROV.wasAttributedTo, self.was_attributed_to.id_),
+        ]
+
+
+@dataclass(frozen=True)
+class Change:
+    agent: Agent
+    activity: Activity
+    entity: Entity
+    description: str
+
+    def as_triples(self):
+        return self.agent.as_triples() + self.activity.as_triples() + self.entity.as_triples()
+
+    @classmethod
+    def record(cls, activity: str, start: datetime, end: datetime, description: str):
+        """User friendly method to record a change that occurred in the graph store."""
+        agent = Agent()
+        activity = Activity(used=activity, was_associated_with=agent, started_at_time=start, ended_at_time=end)
+        entity = Entity(was_generated_by=activity, was_attributed_to=agent)
+        return cls(agent, activity, entity, description)
+
+
+T_Change = TypeVar("T_Change", bound=Change)
+
+
+class Provenance(UserList[T_Change]):
+    def __init__(self, changes: Sequence[T_Change] | None = None):
+        super().__init__(changes or [])
+
+    def did_this_happen(self, this: str) -> bool:
+        return any(change.description == this for change in self)
+
+    def __delitem__(self, *args, **kwargs):
+        raise TypeError("Cannot delete change from provenance")
+
+    def __setitem__(self, *args, **kwargs):
+        raise TypeError("Cannot modify change from provenance")

cognite/neat/issues.py

@@ -0,0 +1,150 @@
+import sys
+import warnings
+from abc import ABC, abstractmethod
+from collections import UserList
+from collections.abc import Sequence
+from dataclasses import dataclass
+from functools import total_ordering
+from typing import Any, ClassVar, TypeVar
+from warnings import WarningMessage
+
+import pandas as pd
+
+if sys.version_info < (3, 11):
+    from exceptiongroup import ExceptionGroup
+    from typing_extensions import Self
+else:
+    from typing import Self
+
+
+@total_ordering
+@dataclass(frozen=True)
+class NeatIssue(ABC):
+    description: ClassVar[str]
+    fix: ClassVar[str]
+
+    def message(self) -> str:
+        """Return a human-readable message for the issue.
+
+        This is the default implementation, which returns the description.
+        It is recommended to override this method in subclasses with a more
+        specific message.
+        """
+        return self.description
+
+    @abstractmethod
+    def dump(self) -> dict[str, Any]:
+        """Return a dictionary representation of the issue."""
+        raise NotImplementedError()
+
+    def __lt__(self, other: "NeatIssue") -> bool:
+        if not isinstance(other, NeatIssue):
+            return NotImplemented
+        return (type(self).__name__, self.message()) < (type(other).__name__, other.message())
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, NeatIssue):
+            return NotImplemented
+        return (type(self).__name__, self.message()) == (type(other).__name__, other.message())
+
+
+@dataclass(frozen=True)
+class NeatError(NeatIssue, ABC):
+    def dump(self) -> dict[str, Any]:
+        return {"errorType": type(self).__name__}
+
+    def as_exception(self) -> ValueError:
+        return ValueError(self.message())
+
+
+@dataclass(frozen=True)
+class NeatWarning(NeatIssue, ABC, UserWarning):
+    def dump(self) -> dict[str, Any]:
+        return {"warningType": type(self).__name__}
+
+    @classmethod
+    def from_warning(cls, warning: WarningMessage) -> "NeatWarning":
+        return DefaultWarning.from_warning_message(warning)
+
+
+@dataclass(frozen=True)
+class DefaultWarning(NeatWarning):
+    description = "A warning was raised during validation."
+    fix = "No fix is available."
+
+    warning: str | Warning
+    category: type[Warning]
+    source: str | None = None
+
+    def dump(self) -> dict[str, Any]:
+        output = super().dump()
+        output["msg"] = str(self.warning)
+        output["category"] = self.category.__name__
+        output["source"] = self.source
+        return output
+
+    @classmethod
+    def from_warning_message(cls, warning: WarningMessage) -> NeatWarning:
+        if isinstance(warning.message, NeatWarning):
+            return warning.message
+
+        return cls(
+            warning=warning.message,
+            category=warning.category,
+            source=warning.source,
+        )
+
+    def message(self) -> str:
+        return str(self.warning)
+
+
+T_NeatIssue = TypeVar("T_NeatIssue", bound=NeatIssue)
+
+
+class NeatIssueList(UserList[T_NeatIssue], ABC):
+    def __init__(self, issues: Sequence[T_NeatIssue] | None = None, title: str | None = None):
+        super().__init__(issues or [])
+        self.title = title
+
+    @property
+    def errors(self) -> Self:
+        return type(self)([issue for issue in self if isinstance(issue, NeatError)])  # type: ignore[misc]
+
+    @property
+    def has_errors(self) -> bool:
+        return any(isinstance(issue, NeatError) for issue in self)
+
+    @property
+    def warnings(self) -> Self:
+        return type(self)([issue for issue in self if isinstance(issue, NeatWarning)])  # type: ignore[misc]
+
+    def as_errors(self) -> ExceptionGroup:
+        return ExceptionGroup(
+            "Operation failed",
+            [ValueError(issue.message()) for issue in self if isinstance(issue, NeatError)],
+        )
+
+    def trigger_warnings(self) -> None:
+        for warning in [issue for issue in self if isinstance(issue, NeatWarning)]:
+            warnings.warn(warning, stacklevel=2)
+
+    def to_pandas(self) -> pd.DataFrame:
+        return pd.DataFrame([issue.dump() for issue in self])
+
+    def _repr_html_(self) -> str | None:
+        return self.to_pandas()._repr_html_()  # type: ignore[operator]
+
+    def as_exception(self) -> "MultiValueError":
+        return MultiValueError(self.errors)
+
+
+class MultiValueError(ValueError):
+    """This is a container for multiple errors.
+
+    It is used in the pydantic field_validator/model_validator to collect multiple errors, which
+    can then be caught in a try-except block and returned as an IssueList.
+
+    """
+
+    def __init__(self, errors: Sequence[T_NeatIssue]):
+        self.errors = list(errors)

cognite/neat/rules/exporters/_base.py

@@ -7,8 +7,7 @@ from cognite.client import CogniteClient
 
 from cognite.neat.rules._shared import Rules
 from cognite.neat.rules.models import DMSRules, InformationRules, RoleTypes
-
-from ._models import UploadResult
+from cognite.neat.utils.upload import UploadDiffsCount
 
 T_Export = TypeVar("T_Export")
 
@@ -38,5 +37,5 @@ class BaseExporter(ABC, Generic[T_Export]):
 
 class CDFExporter(BaseExporter[T_Export]):
     @abstractmethod
-    def export_to_cdf(self, rules: Rules, client: CogniteClient, dry_run: bool = False) -> Iterable[UploadResult]:
+    def export_to_cdf(self, rules: Rules, client: CogniteClient, dry_run: bool = False) -> Iterable[UploadDiffsCount]:
         raise NotImplementedError

cognite/neat/rules/exporters/_rules2dms.py

@@ -31,9 +31,9 @@ from cognite.neat.utils.cdf_loaders import (
     TransformationLoader,
     ViewLoader,
 )
+from cognite.neat.utils.upload import UploadDiffsCount
 
 from ._base import CDFExporter
-from ._models import UploadResult
 
 Component: TypeAlias = Literal["all", "spaces", "data_models", "views", "containers", "node_types"]
 
@@ -123,7 +123,7 @@ class DMSExporter(CDFExporter[DMSSchema]):
             raise ValueError(f"{type(rules).__name__} cannot be exported to DMS")
         return dms_rules.as_schema(include_pipeline=self.export_pipeline, instance_space=self.instance_space)
 
-    def delete_from_cdf(self, rules: Rules, client: CogniteClient, dry_run: bool = False) -> Iterable[UploadResult]:
+    def delete_from_cdf(self, rules: Rules, client: CogniteClient, dry_run: bool = False) -> Iterable[UploadDiffsCount]:
         schema, to_export = self._prepare_schema_and_exporters(rules, client)
 
         # we need to reverse order in which we are picking up the items to delete
@@ -161,7 +161,7 @@ class DMSExporter(CDFExporter[DMSSchema]):
                         deleted -= failed_deleted
                         error_messages.append(f"Failed delete: {e.message}")
 
-            yield UploadResult(
+            yield UploadDiffsCount(
                 name=loader.resource_name,
                 deleted=deleted,
                 skipped=0,
@@ -169,7 +169,7 @@ class DMSExporter(CDFExporter[DMSSchema]):
                 error_messages=error_messages,
             )
 
-    def export_to_cdf(self, rules: Rules, client: CogniteClient, dry_run: bool = False) -> Iterable[UploadResult]:
+    def export_to_cdf(self, rules: Rules, client: CogniteClient, dry_run: bool = False) -> Iterable[UploadDiffsCount]:
         schema, to_export = self._prepare_schema_and_exporters(rules, client)
 
         # The conversion from DMS to GraphQL does not seem to be triggered even if the views
@@ -250,7 +250,7 @@ class DMSExporter(CDFExporter[DMSSchema]):
                         changed -= failed_changed
                         error_messages.append(e.message)
 
-            yield UploadResult(
+            yield UploadDiffsCount(
                 name=loader.resource_name,
                 created=created,
                 changed=changed,

cognite/neat/rules/importers/_base.py

@@ -118,4 +118,4 @@ def _handle_issues(
             future_result._result = "success"
         finally:
             if warning_logger:
-                issues.extend([warning_cls.from_warning(warning) for warning in warning_logger])
+                issues.extend([warning_cls.from_warning(warning) for warning in warning_logger])  # type: ignore[misc]