cognite-toolkit 0.6.97__py3-none-any.whl → 0.7.39__py3-none-any.whl

cognite_toolkit/_cdf_tk/commands/_migrate/migration_io.py

@@ -1,29 +1,59 @@
-from collections.abc import Iterator, Mapping, Sequence
-from typing import ClassVar, cast
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+from typing import ClassVar, Literal, cast
 
-from cognite.client.data_classes.data_modeling import InstanceApply, NodeId
+from cognite.client.data_classes import Annotation
+from cognite.client.data_classes.data_modeling import EdgeId, InstanceApply, NodeId
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
-from cognite_toolkit._cdf_tk.client.data_classes.pending_instances_ids import PendingInstanceId
+from cognite_toolkit._cdf_tk.client.data_classes.legacy.pending_instances_ids import PendingInstanceId
+from cognite_toolkit._cdf_tk.client.data_classes.three_d import (
+    AssetMappingDMRequest,
+    AssetMappingResponse,
+    ThreeDModelResponse,
+)
+from cognite_toolkit._cdf_tk.commands._migrate.data_classes import ThreeDMigrationRequest
 from cognite_toolkit._cdf_tk.constants import MISSING_EXTERNAL_ID, MISSING_INSTANCE_SPACE
-from cognite_toolkit._cdf_tk.exceptions import ToolkitNotImplementedError
+from cognite_toolkit._cdf_tk.exceptions import ToolkitNotImplementedError, ToolkitValueError
 from cognite_toolkit._cdf_tk.storageio import (
+    AnnotationIO,
     HierarchyIO,
     InstanceIO,
+    T_Selector,
     UploadableStorageIO,
 )
 from cognite_toolkit._cdf_tk.storageio._base import Page, UploadItem
+from cognite_toolkit._cdf_tk.storageio.selectors import (
+    ThreeDModelFilteredSelector,
+    ThreeDModelIdSelector,
+    ThreeDSelector,
+)
+from cognite_toolkit._cdf_tk.tk_warnings import MediumSeverityWarning
 from cognite_toolkit._cdf_tk.utils.collection import chunker_sequence
-from cognite_toolkit._cdf_tk.utils.http_client import HTTPClient, HTTPMessage, ItemsRequest, SuccessResponseItems
+from cognite_toolkit._cdf_tk.utils.http_client import (
+    FailedResponse,
+    HTTPClient,
+    HTTPMessage,
+    ItemsRequest,
+    SimpleBodyRequest,
+    SuccessResponseItems,
+    ToolkitAPIError,
+)
 from cognite_toolkit._cdf_tk.utils.useful_types import (
-    AssetCentricKind,
+    AssetCentricKindExtended,
     AssetCentricType,
     JsonVal,
     T_AssetCentricResource,
 )
 
-from .data_classes import AssetCentricMapping, AssetCentricMappingList, MigrationMapping, MigrationMappingList
+from .data_classes import (
+    AnnotationMapping,
+    AssetCentricMapping,
+    AssetCentricMappingList,
+    MigrationMapping,
+    MigrationMappingList,
+)
 from .data_model import INSTANCE_SOURCE_VIEW_ID
+from .default_mappings import ASSET_ANNOTATIONS_ID, FILE_ANNOTATIONS_ID
 from .selectors import AssetCentricMigrationSelector, MigrateDataSetSelector, MigrationCSVFileSelector
 
 
@@ -37,7 +67,7 @@ class AssetCentricMigrationIO(
     CHUNK_SIZE = 1000
     UPLOAD_ENDPOINT = InstanceIO.UPLOAD_ENDPOINT
 
-    PENDING_INSTANCE_ID_ENDPOINT_BY_KIND: ClassVar[Mapping[AssetCentricKind, str]] = {
+    PENDING_INSTANCE_ID_ENDPOINT_BY_KIND: ClassVar[Mapping[AssetCentricKindExtended, str]] = {
         "TimeSeries": "/timeseries/set-pending-instance-ids",
         "FileMetadata": "/files/set-pending-instance-ids",
     }
@@ -113,8 +143,8 @@ class AssetCentricMigrationIO(
             yield mapping_list
 
     @staticmethod
-    def _kind_to_resource_type(kind: AssetCentricKind) -> AssetCentricType:
-        mapping: dict[AssetCentricKind, AssetCentricType] = {
+    def _kind_to_resource_type(kind: AssetCentricKindExtended) -> AssetCentricType:
+        mapping: dict[AssetCentricKindExtended, AssetCentricType] = {
             "Assets": "asset",
             "Events": "event",
             "TimeSeries": "timeseries",
@@ -199,3 +229,339 @@ class AssetCentricMigrationIO(
             pending_instance_id=NodeId(item.space, item.external_id),
             id=id_,
         )
+
+
+class AnnotationMigrationIO(
+    UploadableStorageIO[AssetCentricMigrationSelector, AssetCentricMapping[Annotation], InstanceApply]
+):
+    """IO class for migrating Annotations.
+
+    Args:
+        client: The ToolkitClient to use for CDF interactions.
+        instance_space: The instance space to use for the migrated annotations.
+        default_asset_annotation_mapping: The default ingestion mapping to use for asset-linked annotations.
+        default_file_annotation_mapping: The default ingestion mappingto use for file-linked annotations.
+
+    """
+
+    KIND = "AnnotationMigration"
+    SUPPORTED_DOWNLOAD_FORMATS = frozenset({".parquet", ".csv", ".ndjson"})
+    SUPPORTED_COMPRESSIONS = frozenset({".gz"})
+    SUPPORTED_READ_FORMATS = frozenset({".parquet", ".csv", ".ndjson", ".yaml", ".yml"})
+    CHUNK_SIZE = 1000
+    UPLOAD_ENDPOINT = InstanceIO.UPLOAD_ENDPOINT
+
+    SUPPORTED_ANNOTATION_TYPES = frozenset({"diagrams.AssetLink", "diagrams.FileLink"})
+
+    def __init__(
+        self,
+        client: ToolkitClient,
+        instance_space: str | None = None,
+        default_asset_annotation_mapping: str | None = None,
+        default_file_annotation_mapping: str | None = None,
+    ) -> None:
+        super().__init__(client)
+        self.annotation_io = AnnotationIO(client)
+        self.instance_space = instance_space
+        self.default_asset_annotation_mapping = default_asset_annotation_mapping or ASSET_ANNOTATIONS_ID
+        self.default_file_annotation_mapping = default_file_annotation_mapping or FILE_ANNOTATIONS_ID
+
+    def as_id(self, item: AssetCentricMapping[Annotation]) -> str:
+        return f"Annotation_{item.mapping.id}"
+
+    def count(self, selector: AssetCentricMigrationSelector) -> int | None:
+        if isinstance(selector, MigrationCSVFileSelector):
+            return len(selector.items)
+        else:
+            # There is no efficient way to count annotations in CDF.
+            return None
+
+    def stream_data(self, selector: AssetCentricMigrationSelector, limit: int | None = None) -> Iterable[Page]:
+        if isinstance(selector, MigrateDataSetSelector):
+            iterator = self._stream_from_dataset(selector, limit)
+        elif isinstance(selector, MigrationCSVFileSelector):
+            iterator = self._stream_from_csv(selector, limit)
+        else:
+            raise ToolkitNotImplementedError(f"Selector {type(selector)} is not supported for stream_data")
+        yield from (Page(worker_id="main", items=items) for items in iterator)
+
+    def _stream_from_dataset(
+        self, selector: MigrateDataSetSelector, limit: int | None = None
+    ) -> Iterator[Sequence[AssetCentricMapping[Annotation]]]:
+        if self.instance_space is None:
+            raise ToolkitValueError("Instance space must be provided for dataset-based annotation migration.")
+        asset_centric_selector = selector.as_asset_centric_selector()
+        for data_chunk in self.annotation_io.stream_data(asset_centric_selector, limit):
+            mapping_list = AssetCentricMappingList[Annotation]([])
+            for resource in data_chunk.items:
+                if resource.annotation_type not in self.SUPPORTED_ANNOTATION_TYPES:
+                    # This should not happen, as the annotation_io should already filter these out.
+                    # This is just in case.
+                    continue
+                mapping = AnnotationMapping(
+                    instance_id=EdgeId(space=self.instance_space, external_id=f"annotation_{resource.id!r}"),
+                    id=resource.id,
+                    ingestion_view=self._get_mapping(selector.ingestion_mapping, resource),
+                    preferred_consumer_view=selector.preferred_consumer_view,
+                    # The PySDK is poorly typed.
+                    annotation_type=resource.annotation_type,  # type: ignore[arg-type]
+                )
+                mapping_list.append(AssetCentricMapping(mapping=mapping, resource=resource))
+            yield mapping_list
+
+    def _stream_from_csv(
+        self, selector: MigrationCSVFileSelector, limit: int | None = None
+    ) -> Iterator[Sequence[AssetCentricMapping[Annotation]]]:
+        items = selector.items
+        if limit is not None:
+            items = MigrationMappingList(items[:limit])
+        chunk: list[AssetCentricMapping[Annotation]] = []
+        for current_batch in chunker_sequence(items, self.CHUNK_SIZE):
+            resources = self.client.annotations.retrieve_multiple(current_batch.get_ids())
+            resources_by_id = {resource.id: resource for resource in resources}
+            not_found = 0
+            incorrect_type_count = 0
+            for mapping in current_batch:
+                resource = resources_by_id.get(mapping.id)
+                if resource is None:
+                    not_found += 1
+                    continue
+                if resource.annotation_type not in self.SUPPORTED_ANNOTATION_TYPES:
+                    incorrect_type_count += 1
+                    continue
+                mapping.ingestion_view = self._get_mapping(mapping.ingestion_view, resource)
+                chunk.append(AssetCentricMapping(mapping=mapping, resource=resource))
+            if chunk:
+                yield chunk
+                chunk = []
+            if not_found:
+                MediumSeverityWarning(
+                    f"Could not find {not_found} annotations referenced in the CSV file. They will be skipped during migration."
+                ).print_warning(include_timestamp=True, console=self.client.console)
+            if incorrect_type_count:
+                MediumSeverityWarning(
+                    f"Found {incorrect_type_count} annotations with unsupported types. Only 'diagrams.AssetLink' and "
+                    "'diagrams.FileLink' are supported. These annotations will be skipped during migration."
+                ).print_warning(include_timestamp=True, console=self.client.console)
+
+    def _get_mapping(self, current_mapping: str | None, resource: Annotation) -> str:
+        try:
+            return (
+                current_mapping
+                or {
+                    "diagrams.AssetLink": self.default_asset_annotation_mapping,
+                    "diagrams.FileLink": self.default_file_annotation_mapping,
+                }[resource.annotation_type]
+            )
+        except KeyError as e:
+            raise ToolkitValueError(
+                f"Could not determine default ingestion view for annotation type '{resource.annotation_type}'. "
+                "Please specify the ingestion view explicitly in the CSV file."
+            ) from e
+
+    def json_to_resource(self, item_json: dict[str, JsonVal]) -> InstanceApply:
+        raise NotImplementedError("Deserializing Annotation Migrations from JSON is not supported.")
+
+    def data_to_json_chunk(
+        self,
+        data_chunk: Sequence[AssetCentricMapping[Annotation]],
+        selector: AssetCentricMigrationSelector | None = None,
+    ) -> list[dict[str, JsonVal]]:
+        raise NotImplementedError("Serializing Annotation Migrations to JSON is not supported.")
+
+
+class ThreeDMigrationIO(UploadableStorageIO[ThreeDSelector, ThreeDModelResponse, ThreeDMigrationRequest]):
+    """IO class for downloading and migrating 3D models.
+
+    Args:
+        client: The ToolkitClient to use for CDF interactions.
+        data_model_type: The type of 3D data model to download. Either "classic" or "DM".
+
+    """
+
+    KIND = "3DMigration"
+    SUPPORTED_DOWNLOAD_FORMATS = frozenset({".ndjson"})
+    SUPPORTED_COMPRESSIONS = frozenset({".gz"})
+    SUPPORTED_READ_FORMATS = frozenset({".ndjson"})
+    DOWNLOAD_LIMIT = 1000
+    CHUNK_SIZE = 1
+    UPLOAD_ENDPOINT = "/3d/migrate/models"
+    REVISION_ENDPOINT = "/3d/migrate/revisions"
+
+    def __init__(self, client: ToolkitClient, data_model_type: Literal["classic", "data modeling"] = "classic") -> None:
+        super().__init__(client)
+        self.data_model_type = data_model_type
+
+    def as_id(self, item: ThreeDModelResponse) -> str:
+        return f"{item.name}_{item.id!s}"
+
+    def _is_selected(self, item: ThreeDModelResponse, included_models: set[int] | None) -> bool:
+        return self._is_correct_type(item) and (included_models is None or item.id in included_models)
+
+    def _is_correct_type(self, item: ThreeDModelResponse) -> bool:
+        if self.data_model_type == "classic":
+            return item.space is None
+        else:
+            return item.space is not None
+
+    def stream_data(self, selector: ThreeDSelector, limit: int | None = None) -> Iterable[Page[ThreeDModelResponse]]:
+        published: bool | None = None
+        if isinstance(selector, ThreeDModelFilteredSelector):
+            published = selector.published
+        included_models: set[int] | None = None
+        if isinstance(selector, ThreeDModelIdSelector):
+            included_models = set(selector.ids)
+        cursor: str | None = None
+        total = 0
+        while True:
+            request_limit = min(self.DOWNLOAD_LIMIT, limit - total) if limit is not None else self.DOWNLOAD_LIMIT
+            response = self.client.tool.three_d.models.iterate(
+                published=published, include_revision_info=True, limit=request_limit, cursor=cursor
+            )
+            items = [item for item in response.items if self._is_selected(item, included_models)]
+            total += len(items)
+            if items:
+                yield Page(worker_id="main", items=items, next_cursor=response.next_cursor)
+            if response.next_cursor is None:
+                break
+            cursor = response.next_cursor
+
+    def count(self, selector: ThreeDSelector) -> int | None:
+        # There is no efficient way to count 3D models in CDF.
+        return None
+
+    def data_to_json_chunk(
+        self, data_chunk: Sequence[ThreeDModelResponse], selector: ThreeDSelector | None = None
+    ) -> list[dict[str, JsonVal]]:
+        raise NotImplementedError("Deserializing Annotation Migrations from JSON is not supported.")
+
+    def json_to_resource(self, item_json: dict[str, JsonVal]) -> ThreeDMigrationRequest:
+        raise NotImplementedError("Deserializing ThreeD Migrations from JSON is not supported.")
+
+    def upload_items(
+        self,
+        data_chunk: Sequence[UploadItem[ThreeDMigrationRequest]],
+        http_client: HTTPClient,
+        selector: ThreeDSelector | None = None,
+    ) -> Sequence[HTTPMessage]:
+        """Migrate 3D models by uploading them to the migrate/models endpoint."""
+        if len(data_chunk) > self.CHUNK_SIZE:
+            raise RuntimeError(f"Uploading more than {self.CHUNK_SIZE} 3D models at a time is not supported.")
+
+        results: list[HTTPMessage] = []
+        responses = http_client.request_with_retries(
+            message=ItemsRequest(
+                endpoint_url=self.client.config.create_api_url(self.UPLOAD_ENDPOINT),
+                method="POST",
+                items=list(data_chunk),
+            )
+        )
+        if (
+            failed_response := next((res for res in responses if isinstance(res, FailedResponse)), None)
+        ) and failed_response.status_code == 400:
+            raise ToolkitAPIError("3D model migration failed. You need to enable the 3D migration alpha feature flag.")
+
+        results.extend(responses)
+        success_ids = {id for res in responses if isinstance(res, SuccessResponseItems) for id in res.ids}
+        for data in data_chunk:
+            if data.source_id not in success_ids:
+                continue
+            revision = http_client.request_with_retries(
+                message=SimpleBodyRequest(
+                    endpoint_url=self.client.config.create_api_url(self.REVISION_ENDPOINT),
+                    method="POST",
+                    body_content={"items": [data.item.revision.dump(camel_case=True)]},
+                )
+            )
+            results.extend(revision.as_item_responses(data.source_id))
+        return results
+
+
+class ThreeDAssetMappingMigrationIO(UploadableStorageIO[ThreeDSelector, AssetMappingResponse, AssetMappingDMRequest]):
+    KIND = "3DMigrationAssetMapping"
+    SUPPORTED_DOWNLOAD_FORMATS = frozenset({".ndjson"})
+    SUPPORTED_COMPRESSIONS = frozenset({".gz"})
+    SUPPORTED_READ_FORMATS = frozenset({".ndjson"})
+    DOWNLOAD_LIMIT = 1000
+    CHUNK_SIZE = 100
+    UPLOAD_ENDPOINT = "/3d/models/{modelId}/revisions/{revisionId}/mappings"
+
+    def __init__(self, client: ToolkitClient, object_3D_space: str, cad_node_space: str) -> None:
+        super().__init__(client)
+        self.object_3D_space = object_3D_space
+        self.cad_node_space = cad_node_space
+        # We can only migrate asset mappings for 3D models that are already migrated to data modeling.
+        self._3D_io = ThreeDMigrationIO(client, data_model_type="data modeling")
+
+    def as_id(self, item: AssetMappingResponse) -> str:
+        return f"AssetMapping_{item.model_id!s}_{item.revision_id!s}_{item.asset_id!s}"
+
+    def stream_data(self, selector: ThreeDSelector, limit: int | None = None) -> Iterable[Page[AssetMappingResponse]]:
+        total = 0
+        for three_d_page in self._3D_io.stream_data(selector, None):
+            for model in three_d_page.items:
+                if model.last_revision_info is None or model.last_revision_info.revision_id is None:
+                    # No revisions, so no asset mappings to
+                    continue
+                cursor: str | None = None
+                while True:
+                    request_limit = (
+                        min(self.DOWNLOAD_LIMIT, limit - total) if limit is not None else self.DOWNLOAD_LIMIT
+                    )
+                    if limit is not None and total >= limit:
+                        return
+                    response = self.client.tool.three_d.asset_mappings.iterate(
+                        model_id=model.id,
+                        revision_id=model.last_revision_info.revision_id,
+                        cursor=cursor,
+                        limit=request_limit,
+                    )
+                    items = response.items
+                    total += len(items)
+                    if items:
+                        yield Page(worker_id="main", items=items, next_cursor=response.next_cursor)
+                    if response.next_cursor is None:
+                        break
+                    cursor = response.next_cursor
+
+    def count(self, selector: ThreeDSelector) -> int | None:
+        # There is no efficient way to count 3D asset mappings in CDF.
+        return None
+
+    def upload_items(
+        self,
+        data_chunk: Sequence[UploadItem[AssetMappingDMRequest]],
+        http_client: HTTPClient,
+        selector: T_Selector | None = None,
+    ) -> Sequence[HTTPMessage]:
+        """Migrate 3D asset mappings by uploading them to the migrate/asset-mappings endpoint."""
+        if not data_chunk:
+            return []
+        # Assume all items in the chunk belong to the same model and revision, they should
+        # if the .stream_data method is used for downloading.
+        first = data_chunk[0]
+        model_id = first.item.model_id
+        revision_id = first.item.revision_id
+        endpoint = self.UPLOAD_ENDPOINT.format(modelId=model_id, revisionId=revision_id)
+        responses = http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self.client.config.create_api_url(endpoint),
+                method="POST",
+                items=list(data_chunk),
+                extra_body_fields={
+                    "dmsContextualizationConfig": {
+                        "object3DSpace": self.object_3D_space,
+                        "cadNodeSpace": self.cad_node_space,
+                    }
+                },
+            )
+        )
+        return responses
+
+    def json_to_resource(self, item_json: dict[str, JsonVal]) -> AssetMappingDMRequest:
+        raise NotImplementedError("Deserializing 3D Asset Mappings from JSON is not supported.")
+
+    def data_to_json_chunk(
+        self, data_chunk: Sequence[AssetMappingResponse], selector: ThreeDSelector | None = None
+    ) -> list[dict[str, JsonVal]]:
+        raise NotImplementedError("Serializing 3D Asset Mappings to JSON is not supported.")

cognite_toolkit/_cdf_tk/commands/_migrate/selectors.py

@@ -8,11 +8,11 @@ from cognite.client.data_classes.data_modeling import ViewId
 from cognite_toolkit._cdf_tk.commands._migrate.data_classes import MigrationMappingList
 from cognite_toolkit._cdf_tk.storageio import DataSelector
 from cognite_toolkit._cdf_tk.storageio.selectors import DataSetSelector
-from cognite_toolkit._cdf_tk.utils.useful_types import AssetCentricKind
+from cognite_toolkit._cdf_tk.utils.useful_types import AssetCentricKindExtended
 
 
 class AssetCentricMigrationSelector(DataSelector, ABC):
-    kind: AssetCentricKind
+    kind: AssetCentricKindExtended
 
     @abstractmethod
     def get_ingestion_mappings(self) -> list[str]:
@@ -41,7 +41,7 @@ class MigrationCSVFileSelector(AssetCentricMigrationSelector):
 
 class MigrateDataSetSelector(AssetCentricMigrationSelector):
     type: Literal["migrateDataSet"] = "migrateDataSet"
-    kind: AssetCentricKind
+    kind: AssetCentricKindExtended
     data_set_external_id: str
     ingestion_mapping: str | None = None
     preferred_consumer_view: ViewId | None = None
@@ -63,4 +63,9 @@ class MigrateDataSetSelector(AssetCentricMigrationSelector):
         return [self.ingestion_mapping] if self.ingestion_mapping else []
 
     def as_asset_centric_selector(self) -> DataSetSelector:
-        return DataSetSelector(data_set_external_id=self.data_set_external_id, kind=self.kind)
+        if self.kind == "Annotations":
+            # Annotations are connected to file metadata, so we need to download the file metadata
+            # and look up the annotations connected to each file metadata.
+            return DataSetSelector(data_set_external_id=self.data_set_external_id, kind="FileMetadata")
+        else:
+            return DataSetSelector(data_set_external_id=self.data_set_external_id, kind=self.kind)

cognite_toolkit/_cdf_tk/commands/_profile.py

@@ -21,7 +21,7 @@ from rich.spinner import Spinner
 from rich.table import Table
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
-from cognite_toolkit._cdf_tk.client.data_classes.raw import RawProfileResults, RawTable
+from cognite_toolkit._cdf_tk.client.data_classes.legacy.raw import RawProfileResults, RawTable
 from cognite_toolkit._cdf_tk.constants import MAX_ROW_ITERATION_RUN_QUERY
 from cognite_toolkit._cdf_tk.exceptions import ToolkitMissingDependencyError, ToolkitThrottledError, ToolkitValueError
 from cognite_toolkit._cdf_tk.utils.aggregators import (

cognite_toolkit/_cdf_tk/commands/_purge.py

@@ -7,12 +7,7 @@ from typing import Literal, cast
 
 import questionary
 from cognite.client.data_classes import DataSetUpdate
-from cognite.client.data_classes._base import CogniteResourceList
-from cognite.client.data_classes.data_modeling import (
-    EdgeList,
-    NodeId,
-    NodeList,
-)
+from cognite.client.data_classes.data_modeling import Edge, NodeId
 from cognite.client.data_classes.data_modeling.statistics import SpaceStatistics
 from cognite.client.exceptions import CogniteAPIError
 from cognite.client.utils._identifier import InstanceId
@@ -21,6 +16,7 @@ from rich.console import Console
 from rich.panel import Panel
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
+from cognite_toolkit._cdf_tk.client.data_classes.instance_api import TypedInstanceIdentifier
 from cognite_toolkit._cdf_tk.cruds import (
     AssetCRUD,
     ContainerCRUD,
@@ -46,6 +42,7 @@ from cognite_toolkit._cdf_tk.exceptions import (
     AuthorizationError,
     ToolkitMissingResourceError,
 )
+from cognite_toolkit._cdf_tk.protocols import ResourceResponseProtocol
 from cognite_toolkit._cdf_tk.storageio import InstanceIO
 from cognite_toolkit._cdf_tk.storageio.selectors import InstanceSelector
 from cognite_toolkit._cdf_tk.tk_warnings import (
@@ -63,10 +60,10 @@ from cognite_toolkit._cdf_tk.utils.aggregators import (
     TimeSeriesAggregator,
 )
 from cognite_toolkit._cdf_tk.utils.http_client import (
-    FailedRequestItems,
-    FailedResponseItems,
     HTTPClient,
     ItemsRequest,
+    ItemsRequest2,
+    ItemsSuccessResponse2,
     SuccessResponseItems,
 )
 from cognite_toolkit._cdf_tk.utils.producer_worker import ProducerWorkerExecutor
@@ -107,7 +104,7 @@ class ToDelete(ABC):
     @abstractmethod
     def get_process_function(
         self, client: ToolkitClient, console: Console, verbose: bool, process_results: ResourceDeployResult
-    ) -> Callable[[CogniteResourceList], list[JsonVal]]:
+    ) -> Callable[[list[ResourceResponseProtocol]], list[JsonVal]]:
         raise NotImplementedError()
 
     def get_extra_fields(self) -> dict[str, JsonVal]:
@@ -118,9 +115,10 @@ class ToDelete(ABC):
 class DataModelingToDelete(ToDelete):
     def get_process_function(
         self, client: ToolkitClient, console: Console, verbose: bool, process_results: ResourceDeployResult
-    ) -> Callable[[CogniteResourceList], list[JsonVal]]:
-        def as_id(chunk: CogniteResourceList) -> list[JsonVal]:
-            return [item.as_id().dump(include_type=False) for item in chunk]
+    ) -> Callable[[list[ResourceResponseProtocol]], list[JsonVal]]:
+        def as_id(chunk: list[ResourceResponseProtocol]) -> list[JsonVal]:
+            # We know that all data modeling resources implement as_id
+            return [item.as_id().dump(include_type=False) for item in chunk]  # type: ignore[attr-defined]
 
         return as_id
 
@@ -129,11 +127,11 @@ class DataModelingToDelete(ToDelete):
 class EdgeToDelete(ToDelete):
     def get_process_function(
         self, client: ToolkitClient, console: Console, verbose: bool, process_results: ResourceDeployResult
-    ) -> Callable[[CogniteResourceList], list[JsonVal]]:
-        def as_id(chunk: CogniteResourceList) -> list[JsonVal]:
+    ) -> Callable[[list[ResourceResponseProtocol]], list[JsonVal]]:
+        def as_id(chunk: list[ResourceResponseProtocol]) -> list[JsonVal]:
             return [
                 {"space": item.space, "externalId": item.external_id, "instanceType": "edge"}
-                for item in cast(EdgeList, chunk)
+                for item in cast(list[Edge], chunk)
             ]
 
         return as_id
@@ -146,9 +144,10 @@ class NodesToDelete(ToDelete):
 
     def get_process_function(
         self, client: ToolkitClient, console: Console, verbose: bool, process_results: ResourceDeployResult
-    ) -> Callable[[CogniteResourceList], list[JsonVal]]:
-        def check_for_data(chunk: CogniteResourceList) -> list[JsonVal]:
-            node_ids = cast(NodeList, chunk).as_ids()
+    ) -> Callable[[list[ResourceResponseProtocol]], list[JsonVal]]:
+        def check_for_data(chunk: list[ResourceResponseProtocol]) -> list[JsonVal]:
+            # We know that all node resources implement as_id
+            node_ids = [item.as_id() for item in chunk]  # type: ignore[attr-defined]
             found_ids: set[InstanceId] = set()
             if not self.delete_datapoints:
                 timeseries = client.time_series.retrieve_multiple(instance_ids=node_ids, ignore_unknown_ids=True)
@@ -164,8 +163,7 @@ class NodesToDelete(ToDelete):
                 dumped = node_id.dump(include_instance_type=True)
                 # The delete endpoint expects "instanceType" instead of "type"
                 dumped["instanceType"] = dumped.pop("type")
-                # MyPy think complains about invariant here, even though dict[str, str] is a type of JsonVal
-                result.append(dumped)  # type: ignore[arg-type]
+                result.append(dumped)
             return result
 
         return check_for_data
@@ -175,9 +173,10 @@ class NodesToDelete(ToDelete):
 class IdResourceToDelete(ToDelete):
     def get_process_function(
         self, client: ToolkitClient, console: Console, verbose: bool, process_results: ResourceDeployResult
-    ) -> Callable[[CogniteResourceList], list[JsonVal]]:
-        def as_id(chunk: CogniteResourceList) -> list[JsonVal]:
-            return [{"id": item.id} for item in chunk]
+    ) -> Callable[[list[ResourceResponseProtocol]], list[JsonVal]]:
+        def as_id(chunk: list[ResourceResponseProtocol]) -> list[JsonVal]:
+            # We know that all id resources have an id attribute
+            return [{"id": item.id} for item in chunk]  # type: ignore[attr-defined]
 
         return as_id
 
@@ -186,9 +185,10 @@ class IdResourceToDelete(ToDelete):
 class ExternalIdToDelete(ToDelete):
     def get_process_function(
         self, client: ToolkitClient, console: Console, verbose: bool, process_results: ResourceDeployResult
-    ) -> Callable[[CogniteResourceList], list[JsonVal]]:
-        def as_external_id(chunk: CogniteResourceList) -> list[JsonVal]:
-            return [{"externalId": item.external_id} for item in chunk]
+    ) -> Callable[[list[ResourceResponseProtocol]], list[JsonVal]]:
+        def as_external_id(chunk: list[ResourceResponseProtocol]) -> list[JsonVal]:
+            # We know that all external id resources have an external_id attribute
+            return [{"externalId": item.external_id} for item in chunk]  # type: ignore[attr-defined]
 
         return as_external_id
 
@@ -321,7 +321,7 @@ class PurgeCommand(ToolkitCommand):
                     iteration_count = item.total // self.BATCH_SIZE_DM + (
                         1 if item.total % self.BATCH_SIZE_DM > 0 else 0
                     )
-                executor = ProducerWorkerExecutor[CogniteResourceList, list[JsonVal]](
+                executor = ProducerWorkerExecutor[list[ResourceResponseProtocol], list[JsonVal]](
                     download_iterable=self._iterate_batch(
                         item.crud, space, data_set_external_id, batch_size=self.BATCH_SIZE_DM
                     ),
@@ -348,13 +348,13 @@ class PurgeCommand(ToolkitCommand):
     @staticmethod
     def _iterate_batch(
         crud: ResourceCRUD, selected_space: str | None, data_set_external_id: str | None, batch_size: int
-    ) -> Iterable[CogniteResourceList]:
-        batch = crud.list_cls([])
+    ) -> Iterable[list[ResourceResponseProtocol]]:
+        batch: list[ResourceResponseProtocol] = []
         for resource in crud.iterate(space=selected_space, data_set_external_id=data_set_external_id):
             batch.append(resource)
             if len(batch) >= batch_size:
                 yield batch
-                batch = crud.list_cls([])
+                batch = []
         if batch:
             yield batch
 
@@ -715,21 +715,18 @@ class PurgeCommand(ToolkitCommand):
             results.deleted += len(items)
             return
 
-        responses = delete_client.request_with_retries(
-            ItemsRequest(
-                delete_client.config.create_api_url("/models/instances/delete"),
+        responses = delete_client.request_items_retries(
+            ItemsRequest2(
+                endpoint_url=delete_client.config.create_api_url("/models/instances/delete"),
                 method="POST",
-                # MyPy does not understand that InstanceId.load handles dict[str, JsonVal]
-                items=[DeleteItem(item=item, as_id_fun=InstanceId.load) for item in items],  # type: ignore[arg-type]
+                items=[TypedInstanceIdentifier._load(item) for item in items],
             )
         )
         for response in responses:
-            if isinstance(response, SuccessResponseItems):
+            if isinstance(response, ItemsSuccessResponse2):
                 results.deleted += len(response.ids)
-            elif isinstance(response, FailedResponseItems | FailedRequestItems):
-                results.failed += len(response.ids)
             else:
-                results.failed += len(items)
+                results.failed += len(response.ids)
 
     @staticmethod
     def _unlink_timeseries(

cognite_toolkit/_cdf_tk/commands/_questionary_style.py

@@ -0,0 +1,16 @@
+import questionary
+
+custom_style_fancy = questionary.Style(
+    [
+        ("qmark", "fg:#673ab7"),  # token in front of the question
+        ("question", "bold"),  # question text
+        ("answer", "fg:#f44336 bold"),  # submitted answer text behind the question
+        ("pointer", "fg:#673ab7 bold"),  # pointer used in select and checkbox prompts
+        ("highlighted", "fg:#673ab7 bold"),  # pointed-at choice in select and checkbox prompts
+        ("selected", "fg:#673ab7"),  # style for a selected item of a checkbox
+        ("separator", "fg:#cc5454"),  # separator in lists
+        ("instruction", ""),  # user instructions for select, rawselect, checkbox
+        ("text", ""),  # plain text
+        ("disabled", "fg:#858585 italic"),  # disabled choices for select and checkbox prompts
+    ]
+)