Add global key support for Project move_data_rows_to_task_queue (#1284)

Author: vbrodsky

labelbox/__init__.py

@@ -33,3 +33,4 @@
 from labelbox.schema.slice import Slice, CatalogSlice, ModelSlice
 from labelbox.schema.queue_mode import QueueMode
 from labelbox.schema.task_queue import TaskQueue
+from labelbox.schema.identifiables import UniqueIds, GlobalKeys

labelbox/schema/__init__.py

@@ -21,3 +21,4 @@
 import labelbox.schema.batch
 import labelbox.schema.iam_integration
 import labelbox.schema.media_type
+import labelbox.schema.identifiables

labelbox/schema/identifiable.py

@@ -0,0 +1,38 @@
+from abc import ABC, abstractmethod
+from typing import List, Union
+
+
+class Identifiable(ABC):
+    """
+    Base class for any object representing a unique identifier.
+    """
+
+    def __init__(self, key: str):
+        self._key = key
+
+    @property
+    def key(self):
+        return self.key
+
+    def __eq__(self, other):
+        return other.key == self.key
+
+    def __hash__(self):
+        hash(self.key)
+
+    def __str__(self):
+        return self.key.__str__()
+
+
+class UniqueId(Identifiable):
+    """
+    Represents a unique, internally generated id.
+    """
+    pass
+
+
+class GlobalKey(Identifiable):
+    """
+    Represents a user generated id.
+    """
+    pass

labelbox/schema/identifiables.py

@@ -0,0 +1,50 @@
+from enum import Enum
+from typing import List, Union
+
+
+class IdType(str, Enum):
+    """
+    The type of id used to identify a data row.
+    Currently supported types are:
+        - DataRowId: The id assigned to a data row by Labelbox.
+        - GlobalKey: The id assigned to a data row by the user.
+    """
+    DataRowId = "ID"
+    GlobalKey = "GKEY"
+
+
+class Identifiables:
+
+    def __init__(self, iterable, id_type: IdType):
+        """
+        Args:
+            iterable: Iterable of ids (unique or global keys)
+            id_type: The type of id used to identify a data row.
+        """
+        self._iterable = iterable
+        self._index = 0
+        self._id_type = id_type
+
+    def __iter__(self):
+        return iter(self._iterable)
+
+
+class UniqueIds(Identifiables):
+    """
+    Represents a collection of unique, internally generated ids.
+    """
+
+    def __init__(self, iterable: List[str]):
+        super().__init__(iterable, IdType.DataRowId)
+
+
+class GlobalKeys(Identifiables):
+    """
+    Represents a collection of user generated ids.
+    """
+
+    def __init__(self, iterable: List[str]):
+        super().__init__(iterable, IdType.GlobalKey)
+
+
+DataRowIdentifiers = Union[UniqueIds, GlobalKeys]

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,23 +28,16 @@
 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.identifiables import DataRowIdentifiers, UniqueIds
 from labelbox.schema.media_type import MediaType
 from labelbox.schema.queue_mode import QueueMode
 from labelbox.schema.resource_tag import ResourceTag
 from labelbox.schema.task import Task
 from labelbox.schema.task_queue import TaskQueue
-from labelbox.schema.user import User
 
 if TYPE_CHECKING:
     from labelbox import BulkImportRequest
 
-try:
-    datetime.fromisoformat  # type: ignore[attr-defined]
-except AttributeError:
-    from backports.datetime_fromisoformat import MonkeyPatch
-
-    MonkeyPatch.patch_fromisoformat()
-
 try:
     from labelbox.data.serialization import LBV1Converter
 except ImportError:
@@ -643,7 +639,7 @@ def labeler_performance(self) -> PaginatedCollection:
         def create_labeler_performance(client, result):
             result["user"] = Entity.User(client, result["user"])
             # python isoformat doesn't accept Z as utc timezone
-            result["lastActivityTime"] = datetime.fromisoformat(
+            result["lastActivityTime"] = utils.format_iso_from_string(
                 result["lastActivityTime"].replace('Z', '+00:00'))
             return LabelerPerformance(**{
                 utils.snake_case(key): value for key, value in result.items()
@@ -1382,29 +1378,44 @@ 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. A DataIdentifier object can be a UniqueIds or GlobalKeys class.
             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 = UniqueIds(data_row_ids)
+            warnings.warn("Using data row ids will be deprecated. Please use "
+                          "UniqueIds or GlobalKeys instead.")
+
         method = "createBulkAddRowsToQueueTask"
         query_str = """mutation AddDataRowsToTaskQueueAsyncPyApi(
           $projectId: ID!
           $queueId: ID
-          $dataRowIds: [ID!]!
+          $dataRowIdentifiers: AddRowsToTaskQueueViaDataRowIdentifiersInput!
         ) {
           project(where: { id: $projectId }) {
             %s(
-              data: { queueId: $queueId, dataRowIds: $dataRowIds }
+              data: { queueId: $queueId, dataRowIdentifiers: $dataRowIdentifiers }
             ) {
               taskId
             }
@@ -1416,7 +1427,10 @@ 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
+                "dataRowIdentifiers": {
+                    "ids": [id for id in data_row_ids],
+                    "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.identifiables import GlobalKeys, UniqueIds
 
 
 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,25 @@ 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, _, 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)
+
+    review_queue = next(
+        tq for tq in task_queues if tq.queue_type == "MANUAL_REVIEW_QUEUE")
+    project.move_data_rows_to_task_queue(UniqueIds([data_row.uid]),
+                                         review_queue.uid)
+    _validate_moved(project, "MANUAL_REVIEW_QUEUE", 1)

tests/unit/test_unit_identifiables.py

@@ -0,0 +1,15 @@
+from labelbox.schema.identifiables import GlobalKeys, UniqueIds
+
+
+def test_unique_ids():
+    ids = ["a", "b", "c"]
+    identifiables = UniqueIds(ids)
+    assert [i for i in identifiables] == ids
+    assert identifiables._id_type == "ID"
+
+
+def test_global_keys():
+    ids = ["a", "b", "c"]
+    identifiables = GlobalKeys(ids)
+    assert [i for i in identifiables] == ids
+    assert identifiables._id_type == "GKEY"