Update to support identifiable as type in graphql and other identifiable improvements

Author: Val Brodsky

labelbox/schema/identifiable.py

@@ -1,21 +1,20 @@
-from abc import ABC
+from abc import ABC, abstractmethod
 from enum import Enum
 from typing import List, Union
-from typing_extensions import TypeAlias
 
 
 class IdType(str, Enum):
-    DataRowId = "DATA_ROW_ID"
+    DataRowId = "UID"
     GlobalKey = "GLOBAL_KEY"
 
 
 class Identifiable(ABC):
 
-    def __init__(self, keys: Union[str, List[str]]):
+    def __init__(self, keys: Union[str, List[str]], id_type: IdType):
         self._keys = keys
         if isinstance(keys, str):
-            self.keys = [keys]
-        self._id_type = IdType.DataRowId
+            self._keys = [keys]
+        self._id_type = id_type
 
     @property
     def keys(self):
@@ -25,11 +24,12 @@ def keys(self):
     def keys(self, keys):
         self._keys = keys
         if isinstance(keys, str):
-            self.keys = [keys]
+            self._keys = [keys]
 
-    @property
-    def id_type(self):
-        return self._id_type
+    @classmethod
+    @abstractmethod
+    def strings_to_identifiable(cls, keys: Union[str, List[str]]):
+        pass
 
     def __eq__(self, other):
         return other.keys == self.keys
@@ -43,21 +43,23 @@ def __str__(self):
 
 class UniqueIds(Identifiable):
 
+    @classmethod
+    def strings_to_identifiable(cls, keys: Union[str, List[str]]):
+        return cls(keys)
+
     def __init__(self, keys: Union[str, List[str]]):
-        super().__init__(keys)
-        self._id_type = IdType.DataRowId
+        super().__init__(keys, IdType.DataRowId)
 
 
 class GlobalKeys(Identifiable):
 
-    def __init__(self, keys: Union[str, List[str]]):
-        super().__init__(keys)
-        self._id_type = IdType.GlobalKey
-
+    @classmethod
+    def strings_to_identifiable(cls, keys: Union[str, List[str]]):
+        return cls(keys)
 
-DefaultIdentifiable: TypeAlias = UniqueIds
-DataRowIdentifiers: TypeAlias = Union[UniqueIds, GlobalKeys]
+    def __init__(self, keys: Union[str, List[str]]):
+        super().__init__(keys, IdType.GlobalKey)
 
 
-def strings_to_identifiable(keys: Union[str, List[str]]) -> DefaultIdentifiable:
-    return DefaultIdentifiable(keys)
+DefaultIdentifiable = UniqueIds
+DataRowIdentifiers = Union[UniqueIds, GlobalKeys]

labelbox/schema/project.py

@@ -28,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, strings_to_identifiable
+from labelbox.schema.identifiable import DataRowIdentifiers, UniqueIds
 from labelbox.schema.media_type import MediaType
 from labelbox.schema.queue_mode import QueueMode
 from labelbox.schema.resource_tag import ResourceTag
@@ -1403,18 +1403,17 @@ def move_data_rows_to_task_queue(self, data_row_ids, task_queue_id: str):
 
         """
         if isinstance(data_row_ids, list):
-            data_row_ids = strings_to_identifiable(data_row_ids)
+            data_row_ids = UniqueIds.strings_to_identifiable(data_row_ids)
 
         method = "createBulkAddRowsToQueueTask"
         query_str = """mutation AddDataRowsToTaskQueueAsyncPyApi(
           $projectId: ID!
           $queueId: ID
-          $dataRowIds: [ID!]!
-          $idType: IdType!
+          $dataRowIdentifiers: AddRowsToTaskQueueViaDataRowIdentifiersInput!
         ) {
           project(where: { id: $projectId }) {
             %s(
-              data: { queueId: $queueId, dataRowIds: $dataRowIds, idType: $idType }
+              data: { queueId: $queueId, dataRowIdentifiers: $dataRowIdentifiers }
             ) {
               taskId
             }
@@ -1426,8 +1425,10 @@ def move_data_rows_to_task_queue(self, data_row_ids, task_queue_id: str):
             query_str, {
                 "projectId": self.uid,
                 "queueId": task_queue_id,
-                "dataRowIds": data_row_ids.keys,
-                "idType": data_row_ids.id_type,
+                "dataRowIdentifiers": {
+                    "ids": data_row_ids.keys,
+                    "idType": data_row_ids._id_type,
+                },
             },
             timeout=180.0,
             experimental=True)["project"][method]["taskId"]

tests/unit/test_unit_identifiable.py

@@ -1,30 +1,46 @@
-from labelbox.schema.identifiable import GlobalKeys, UniqueIds, strings_to_identifiable
+from labelbox.schema.identifiable import GlobalKeys, UniqueIds
 
 
 def test_unique_ids():
     ids = ["a", "b", "c"]
     identifiable = UniqueIds(ids)
     assert identifiable.keys == ids
-    assert identifiable.id_type == "DATA_ROW_ID"
+    assert identifiable._id_type == "UID"
 
 
 def test_global_keys():
     ids = ["a", "b", "c"]
     identifiable = GlobalKeys(ids)
     assert identifiable.keys == ids
-    assert identifiable.id_type == "GLOBAL_KEY"
+    assert identifiable._id_type == "GLOBAL_KEY"
 
 
-def test_strings_to_identifiable():
+def test_strings_to_identifiable_unique_id():
     ids = ["a", "b", "c"]
-    identifiable = strings_to_identifiable(ids)
+    identifiable = UniqueIds.strings_to_identifiable(ids)
+    assert type(identifiable) == UniqueIds
     assert identifiable.keys == ids
-    assert identifiable.id_type == "DATA_ROW_ID"
+    assert identifiable._id_type == "UID"
 
     ids = "a"
-    identifiable = strings_to_identifiable(ids)
+    identifiable = UniqueIds.strings_to_identifiable(ids)
+    assert type(identifiable) == UniqueIds
     assert identifiable.keys == [ids]
-    assert identifiable.id_type == "DATA_ROW_ID"
+    assert identifiable._id_type == "UID"
+
+
+def test_strings_to_identifiable_global_key():
+    ids = ["a", "b", "c"]
+    identifiable = GlobalKeys.strings_to_identifiable(ids)
+    assert type(identifiable) == GlobalKeys
+    assert identifiable.keys == ids
+    assert identifiable._id_type == "GLOBAL_KEY"
+
+    ids = "a"
+    identifiable = GlobalKeys.strings_to_identifiable(ids)
+    assert type(identifiable) == GlobalKeys
+    assert identifiable.keys == [ids]
+    assert identifiable._id_type == "GLOBAL_KEY"
 
 
 def test__str__():