Implement global key support for move_data_rows_to_task_queue

Author: Val Brodsky

labelbox/schema/project.py

@@ -5,16 +5,19 @@
 from collections import namedtuple
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, TypeVar, Union, overload
 from urllib.parse import urlparse
 
 import requests
 
 from labelbox import parser
 from labelbox import utils
-from labelbox.exceptions import (InvalidQueryError, LabelboxError,
-                                 ProcessingWaitTimeout, ResourceConflict,
-                                 ResourceNotFoundError)
+from labelbox.exceptions import (
+    InvalidQueryError,
+    LabelboxError,
+    ProcessingWaitTimeout,
+    ResourceConflict,
+)
 from labelbox.orm import query
 from labelbox.orm.db_object import DbObject, Deletable, Updateable, experimental
 from labelbox.orm.model import Entity, Field, Relationship
@@ -25,7 +28,7 @@
 from labelbox.schema.export_filters import ProjectExportFilters, validate_datetime, build_filters
 from labelbox.schema.export_params import ProjectExportParams
 from labelbox.schema.export_task import ExportTask
-from labelbox.schema.identifiable import DataRowIdentifiers, IdType, UniqueIds, strings_to_identifiable
+from labelbox.schema.identifiable import DataRowIdentifiers, strings_to_identifiable
 from labelbox.schema.media_type import MediaType
 from labelbox.schema.queue_mode import QueueMode
 from labelbox.schema.resource_tag import ResourceTag
@@ -1375,29 +1378,43 @@ def task_queues(self) -> List[TaskQueue]:
             for field_values in task_queue_values
         ]
 
+    @overload
+    def move_data_rows_to_task_queue(self, data_row_ids: DataRowIdentifiers,
+                                     task_queue_id: str):
+        pass
+
+    @overload
     def move_data_rows_to_task_queue(self, data_row_ids: List[str],
                                      task_queue_id: str):
+        pass
+
+    def move_data_rows_to_task_queue(self, data_row_ids, task_queue_id: str):
         """
 
         Moves data rows to the specified task queue.
 
         Args:
-            data_row_ids: a list of data row ids to be moved
+            data_row_ids: a list of data row ids to be moved. This can be a list of strings or a DataRowIdentifiers object 
+                DataRowIdentifier objects are lists of ids or global keys
             task_queue_id: the task queue id to be moved to, or None to specify the "Done" queue
 
         Returns:
             None if successful, or a raised error on failure
 
         """
+        if isinstance(data_row_ids, list):
+            data_row_ids = strings_to_identifiable(data_row_ids)
+
         method = "createBulkAddRowsToQueueTask"
         query_str = """mutation AddDataRowsToTaskQueueAsyncPyApi(
           $projectId: ID!
           $queueId: ID
           $dataRowIds: [ID!]!
+          $idType: IdType!
         ) {
           project(where: { id: $projectId }) {
             %s(
-              data: { queueId: $queueId, dataRowIds: $dataRowIds }
+              data: { queueId: $queueId, dataRowIds: $dataRowIds, idType: $idType }
             ) {
               taskId
             }
@@ -1409,7 +1426,8 @@ def move_data_rows_to_task_queue(self, data_row_ids: List[str],
             query_str, {
                 "projectId": self.uid,
                 "queueId": task_queue_id,
-                "dataRowIds": data_row_ids
+                "dataRowIds": data_row_ids.keys,
+                "idType": data_row_ids.id_type,
             },
             timeout=180.0,
             experimental=True)["project"][method]["taskId"]

tests/integration/test_task_queue.py

@@ -1,6 +1,7 @@
 import time
 
 from labelbox import Project
+from labelbox.schema.identifiable import GlobalKeys
 
 
 def test_get_task_queue(project: Project):
@@ -11,22 +12,15 @@ def test_get_task_queue(project: Project):
     assert review_queue
 
 
-def test_move_to_task(configured_batch_project_with_label: Project):
-    project, _, data_row, label = configured_batch_project_with_label
-    task_queues = project.task_queues()
-
-    review_queue = next(
-        tq for tq in task_queues if tq.queue_type == "MANUAL_REVIEW_QUEUE")
-    project.move_data_rows_to_task_queue([data_row.uid], review_queue.uid)
-
+def _validate_moved(project, queue_name, data_row_count):
     timeout_seconds = 30
     sleep_time = 2
     while True:
         task_queues = project.task_queues()
         review_queue = next(
-            tq for tq in task_queues if tq.queue_type == "MANUAL_REVIEW_QUEUE")
+            tq for tq in task_queues if tq.queue_type == queue_name)
 
-        if review_queue.data_row_count == 1:
+        if review_queue.data_row_count == data_row_count:
             break
 
         if timeout_seconds <= 0:
@@ -35,3 +29,19 @@ def test_move_to_task(configured_batch_project_with_label: Project):
 
         timeout_seconds -= sleep_time
         time.sleep(sleep_time)
+
+
+def test_move_to_task(configured_batch_project_with_label: Project):
+    project, _, data_row, _ = configured_batch_project_with_label
+    task_queues = project.task_queues()
+
+    review_queue = next(
+        tq for tq in task_queues if tq.queue_type == "MANUAL_REVIEW_QUEUE")
+    project.move_data_rows_to_task_queue([data_row.uid], review_queue.uid)
+    _validate_moved(project, "MANUAL_REVIEW_QUEUE", 1)
+
+    review_queue = next(
+        tq for tq in task_queues if tq.queue_type == "MANUAL_REWORK_QUEUE")
+    project.move_data_rows_to_task_queue(GlobalKeys([data_row.global_key]),
+                                         review_queue.uid)
+    _validate_moved(project, "MANUAL_REWORK_QUEUE", 1)