Add dataset.delete_scene (#199)

* Add dataset.delete_scene

* Improve docstring of delete_scene

Author: gatli

nucleus/data_transfer_object/scenes_list.py

@@ -0,0 +1,18 @@
+from typing import List, Optional
+
+from nucleus.pydantic_base import DictCompatibleImmutableModel
+
+
+class ScenesListEntry(DictCompatibleImmutableModel):
+    """/dataset/{dataset_id}/scenes_list nested payload"""
+
+    id: str
+    reference_id: str
+    type: str
+    metadata: Optional[dict] = {}
+
+
+class ScenesList(DictCompatibleImmutableModel):
+    """/dataset/{dataset_id}/scenes_list payload"""
+
+    scenes: List[ScenesListEntry]

nucleus/dataset.py

@@ -50,6 +50,7 @@
 )
 from .data_transfer_object.dataset_info import DatasetInfo
 from .data_transfer_object.dataset_size import DatasetSize
+from .data_transfer_object.scenes_list import ScenesList, ScenesListEntry
 from .dataset_item import (
     DatasetItem,
     check_all_paths_remote,
@@ -186,13 +187,14 @@ def items(self) -> List[DatasetItem]:
         return constructed_dataset_items
 
     @property
-    def scenes(self) -> List[Dict[str, Any]]:
+    def scenes(self) -> List[ScenesListEntry]:
         """List of ID, reference ID, type, and metadata for all scenes in the Dataset."""
         response = self._client.make_request(
             {}, f"dataset/{self.id}/scenes_list", requests.get
         )
 
-        return response.get("scenes", None)
+        scenes_list = ScenesList.parse_obj(response)
+        return scenes_list.scenes
 
     @sanitize_string_args
     def autotag_items(self, autotag_name, for_scores_greater_than=0):
@@ -738,6 +740,18 @@ def delete_item(self, reference_id: str) -> dict:
             requests.delete,
         )
 
+    @sanitize_string_args
+    def delete_scene(self, reference_id: str):
+        """Deletes a Scene associated with the Dataset
+
+        All items, annotations and predictions associated with the scene will be
+        deleted as well.
+
+        Parameters:
+            reference_id: The user-defined reference ID of the item to delete.
+        """
+        self._client.delete(f"dataset/{self.id}/scene/{reference_id}")
+
     def list_autotags(self):
         """Fetches all autotags of the dataset.
 

nucleus/pydantic_base.py

@@ -14,3 +14,13 @@ class DictCompatibleModel(BaseModel):
 
     def __getitem__(self, key):
         return getattr(self, key)
+
+
+class DictCompatibleImmutableModel(ImmutableModel):
+    """Backwards compatible wrapper where we transform dictionaries into Pydantic Models
+
+    Allows us to access model.key with model["key"].
+    """
+
+    def __getitem__(self, key):
+        return getattr(self, key)

pyproject.toml

@@ -21,7 +21,7 @@ exclude = '''
 
 [tool.poetry]
 name = "scale-nucleus"
-version = "0.5.1"
+version = "0.5.2"
 description = "The official Python client library for Nucleus, the Data Platform for AI"
 license =  "MIT"
 authors = ["Scale AI Nucleus Team <nucleusapi@scaleapi.com>"]

tests/test_scene.py

@@ -1,4 +1,5 @@
 import copy
+import time
 
 import pytest
 
@@ -438,6 +439,31 @@ def test_scene_upload_and_update(dataset_scene):
     }
 
 
+@pytest.mark.integration
+def test_scene_deletion(dataset_scene):
+    payload = TEST_LIDAR_SCENES
+    scenes = [
+        LidarScene.from_json(scene_json) for scene_json in payload[SCENES_KEY]
+    ]
+    update = payload[UPDATE_KEY]
+
+    job = dataset_scene.append(scenes, update=update, asynchronous=True)
+    job.sleep_until_complete()
+
+    uploaded_scenes = dataset_scene.scenes
+    assert len(uploaded_scenes) == len(scenes)
+    assert all(
+        u["reference_id"] == o.reference_id
+        for u, o in zip(uploaded_scenes, scenes)
+    )
+
+    for scene in uploaded_scenes:
+        dataset_scene.delete_scene(scene.reference_id)
+    time.sleep(1)
+    scenes = dataset_scene.scenes
+    assert len(scenes) == 0, f"Expected to delete all scenes, got: {scenes}"
+
+
 @pytest.mark.integration
 def test_scene_upload_async_item_dataset(dataset_item):
     payload = TEST_LIDAR_SCENES