Add Identifiables as iterator

Author: Val Brodsky

labelbox/__init__.py

@@ -33,4 +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.identifiable import UniqueIds, GlobalKeys
+from labelbox.schema.identifiables import UniqueIds, GlobalKeys

labelbox/schema/__init__.py

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

labelbox/schema/identifiable.py

@@ -1,65 +1,29 @@
 from abc import ABC, abstractmethod
-from enum import Enum
 from typing import List, Union
 
 
-class IdType(str, Enum):
-    DataRowId = "UID"
-    GlobalKey = "GLOBAL_KEY"
-
-
 class Identifiable(ABC):
 
-    def __init__(self, keys: Union[str, List[str]], id_type: IdType):
-        self._keys = keys
-        if isinstance(keys, str):
-            self._keys = [keys]
-        self._id_type = id_type
+    def __init__(self, key: str):
+        self._key = key
 
     @property
-    def keys(self):
-        return self._keys
-
-    @keys.setter
-    def keys(self, keys):
-        self._keys = keys
-        if isinstance(keys, str):
-            self._keys = [keys]
-
-    @classmethod
-    @abstractmethod
-    def strings_to_identifiable(cls, keys: Union[str, List[str]]):
-        pass
+    def key(self):
+        return self.key
 
     def __eq__(self, other):
-        return other.keys == self.keys
+        return other.key == self.key
 
     def __hash__(self):
-        hash(self.keys)
+        hash(self.key)
 
     def __str__(self):
-        return self.keys.__str__()
-
-
-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, IdType.DataRowId)
-
-
-class GlobalKeys(Identifiable):
+        return self.key.__str__()
 
-    @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, IdType.GlobalKey)
+class UniqueId(Identifiable):
+    pass
 
 
-DefaultIdentifiable = UniqueIds
-DataRowIdentifiers = Union[UniqueIds, GlobalKeys]
+class GlobalKey(Identifiable):
+    pass

labelbox/schema/identifiables.py

@@ -0,0 +1,33 @@
+from enum import Enum
+from typing import List, Union
+
+
+class IdType(str, Enum):
+    DataRowId = "UID"
+    GlobalKey = "GLOBAL_KEY"
+
+
+class Identifiables:
+
+    def __init__(self, iterable, id_type: IdType):
+        self._iterable = iterable
+        self._index = 0
+        self._id_type = id_type
+
+    def __iter__(self):
+        return iter(self._iterable)
+
+
+class UniqueIds(Identifiables):
+
+    def __init__(self, iterable: List[str]):
+        super().__init__(iterable, IdType.DataRowId)
+
+
+class GlobalKeys(Identifiables):
+
+    def __init__(self, iterable: List[str]):
+        super().__init__(iterable, IdType.GlobalKey)
+
+
+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, UniqueIds
+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
@@ -1403,7 +1403,7 @@ def move_data_rows_to_task_queue(self, data_row_ids, task_queue_id: str):
 
         """
         if isinstance(data_row_ids, list):
-            data_row_ids = UniqueIds.strings_to_identifiable(data_row_ids)
+            data_row_ids = UniqueIds(data_row_ids)
             warnings.warn("Using data row ids will be deprecated. Please use "
                           "UniqueIds or GlobalKeys instead.")
 
@@ -1428,7 +1428,7 @@ def move_data_rows_to_task_queue(self, data_row_ids, task_queue_id: str):
                 "projectId": self.uid,
                 "queueId": task_queue_id,
                 "dataRowIdentifiers": {
-                    "ids": data_row_ids.keys,
+                    "ids": [id for id in data_row_ids],
                     "idType": data_row_ids._id_type,
                 },
             },

tests/integration/test_task_queue.py

@@ -1,7 +1,7 @@
 import time
 
 from labelbox import Project
-from labelbox.schema.identifiable import GlobalKeys
+from labelbox.schema.identifiables import GlobalKeys
 
 
 def test_get_task_queue(project: Project):

tests/unit/test_unit_identifiable.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 == "UID"
+
+
+def test_global_keys():
+    ids = ["a", "b", "c"]
+    identifiables = GlobalKeys(ids)
+    assert [i for i in identifiables] == ids
+    assert identifiables._id_type == "GLOBAL_KEY"