Merge pull request #644 from Labelbox/farkob/batch-management-parity

[AL-2897] Batch management UI parity

Author: jtsodapop

labelbox/schema/batch.py

@@ -68,14 +68,15 @@ def remove_queued_data_rows(self) -> None:
 
         project_id_param = "projectId"
         batch_id_param = "batchId"
-        self.client.execute("""mutation ArchiveBatchPyApi($%s: ID!, $%s: ID!) {
-            project(where: {id: $%s}) { archiveBatch(batchId: $%s) { id archivedAt } }
+        self.client.execute(
+            """mutation RemoveQueuedDataRowsFromBatchPyApi($%s: ID!, $%s: ID!) {
+            project(where: {id: $%s}) { removeQueuedDataRowsFromBatch(batchId: $%s) { id } }
         }""" % (project_id_param, batch_id_param, project_id_param,
                 batch_id_param), {
                     project_id_param: self.project_id,
                     batch_id_param: self.uid
                 },
-                            experimental=True)
+            experimental=True)
 
     def export_data_rows(self,
                          timeout_seconds=120,
@@ -125,3 +126,50 @@ def export_data_rows(self,
             logger.debug("Batch '%s' data row export, waiting for server...",
                          self.uid)
             time.sleep(sleep_time)
+
+    def delete(self) -> None:
+        """ Deletes the given batch.
+
+        Note: Batch deletion for batches that has labels is forbidden. 
+
+        Args:
+            batch (Batch): Batch to remove queued data rows from
+        """
+
+        project_id_param = "projectId"
+        batch_id_param = "batchId"
+        self.client.execute("""mutation DeleteBatchPyApi($%s: ID!, $%s: ID!) {
+            project(where: {id: $%s}) { deleteBatch(batchId: $%s) { deletedBatchId } }
+        }""" % (project_id_param, batch_id_param, project_id_param,
+                batch_id_param), {
+                    project_id_param: self.project_id,
+                    batch_id_param: self.uid
+                },
+                            experimental=True)
+
+    def delete_labels(self, set_labels_as_template=False) -> None:
+        """ Deletes labels that were created for data rows in the batch.
+
+        Args:
+            batch (Batch): Batch to remove queued data rows from
+            set_labels_as_template (bool): When set to true, the deleted labels will be kept as templates.
+        """
+
+        project_id_param = "projectId"
+        batch_id_param = "batchId"
+        type_param = "type"
+        res = self.client.execute(
+            """mutation DeleteBatchLabelsPyApi($%s: ID!, $%s: ID!, $%s: DeleteBatchLabelsType!) {
+            project(where: {id: $%s}) { deleteBatchLabels(batchId: $%s, data:{ type: $%s }) { deletedLabelIds } }
+        }""" % (project_id_param, batch_id_param, type_param, project_id_param,
+                batch_id_param, type_param), {
+                    project_id_param:
+                        self.project_id,
+                    batch_id_param:
+                        self.uid,
+                    type_param:
+                        "RequeueDataWithLabelAsTemplate"
+                        if set_labels_as_template else "RequeueData"
+                },
+            experimental=True)
+        return res

tests/integration/test_batch.py

@@ -50,6 +50,15 @@ def test_archive_batch(configured_project: Project, small_dataset: Dataset):
     assert len(exported_data_rows) == 0
 
 
+def test_delete(configured_project: Project, small_dataset: Dataset):
+    data_rows = [dr.uid for dr in list(small_dataset.export_data_rows())]
+    configured_project.update(queue_mode=Project.QueueMode.Batch)
+    batch = configured_project.create_batch("batch to delete", data_rows)
+    batch.delete()
+
+    assert len(list(configured_project.batches())) == 0
+
+
 def test_batch_project(configured_project: Project, small_dataset: Dataset):
     data_rows = [dr.uid for dr in list(small_dataset.export_data_rows())]
     configured_project.update(queue_mode=Project.QueueMode.Batch)
@@ -80,3 +89,29 @@ def test_export_data_rows(configured_project: Project, dataset: Dataset):
 
     assert len(result) == n_data_rows
     assert set(data_rows) == set(exported_data_rows)
+
+
+@pytest.mark.skip(
+    reason="Test cannot be used effectively with MAL/LabelImport. \
+Fix/Unskip after resolving deletion with MAL/LabelImport")
+def test_delete_labels(configured_project_with_label):
+    project, dataset, _, _ = configured_project_with_label
+
+    data_rows = [dr.uid for dr in list(dataset.export_data_rows())]
+    project.update(queue_mode=Project.QueueMode.Batch)
+    batch = project.create_batch("batch to delete labels", data_rows)
+
+
+@pytest.mark.skip(
+    reason="Test cannot be used effectively with MAL/LabelImport. \
+Fix/Unskip after resolving deletion with MAL/LabelImport")
+def test_delete_labels_with_templates(configured_project: Project,
+                                      small_dataset: Dataset):
+    data_rows = [dr.uid for dr in list(small_dataset.export_data_rows())]
+    configured_project.update(queue_mode=Project.QueueMode.Batch)
+    batch = configured_project.create_batch(
+        "batch to delete labels w templates", data_rows)
+    exported_data_rows = list(batch.export_data_rows())
+    res = batch.delete_labels(labels_as_template=True)
+    exported_data_rows = list(batch.export_data_rows())
+    assert len(exported_data_rows) == 5