Add Identifiable hierarchy

Author: Val Brodsky

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.identifiable 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.identifiable

labelbox/schema/identifiable.py

@@ -0,0 +1,63 @@
+from abc import ABC
+from enum import Enum
+from typing import List, Union
+from typing_extensions import TypeAlias
+
+
+class IdType(str, Enum):
+    DataRowId = "DATA_ROW_ID"
+    GlobalKey = "GLOBAL_KEY"
+
+
+class Identifiable(ABC):
+
+    def __init__(self, keys: Union[str, List[str]]):
+        self._keys = keys
+        if isinstance(keys, str):
+            self.keys = [keys]
+        self._id_type = IdType.DataRowId
+
+    @property
+    def keys(self):
+        return self._keys
+
+    @keys.setter
+    def keys(self, keys):
+        self._keys = keys
+        if isinstance(keys, str):
+            self.keys = [keys]
+
+    @property
+    def id_type(self):
+        return self._id_type
+
+    def __eq__(self, other):
+        return other.keys == self.keys
+
+    def __hash__(self):
+        hash(self.keys)
+
+    def __str__(self):
+        return self.keys.__str__()
+
+
+class UniqueIds(Identifiable):
+
+    def __init__(self, keys: Union[str, List[str]]):
+        super().__init__(keys)
+        self._id_type = IdType.DataRowId
+
+
+class GlobalKeys(Identifiable):
+
+    def __init__(self, keys: Union[str, List[str]]):
+        super().__init__(keys)
+        self._id_type = IdType.GlobalKey
+
+
+DefaultIdentifiable: TypeAlias = UniqueIds
+DataRowIdentifiers: TypeAlias = Union[UniqueIds, GlobalKeys]
+
+
+def strings_to_identifiable(keys: Union[str, List[str]]) -> DefaultIdentifiable:
+    return DefaultIdentifiable(keys)

tests/unit/test_unit_identifiable.py

@@ -0,0 +1,34 @@
+from labelbox.schema.identifiable import GlobalKeys, UniqueIds, strings_to_identifiable
+
+
+def test_unique_ids():
+    ids = ["a", "b", "c"]
+    identifiable = UniqueIds(ids)
+    assert identifiable.keys == ids
+    assert identifiable.id_type == "DATA_ROW_ID"
+
+
+def test_global_keys():
+    ids = ["a", "b", "c"]
+    identifiable = GlobalKeys(ids)
+    assert identifiable.keys == ids
+    assert identifiable.id_type == "GLOBAL_KEY"
+
+
+def test_strings_to_identifiable():
+    ids = ["a", "b", "c"]
+    identifiable = strings_to_identifiable(ids)
+    assert identifiable.keys == ids
+    assert identifiable.id_type == "DATA_ROW_ID"
+
+    ids = "a"
+    identifiable = strings_to_identifiable(ids)
+    assert identifiable.keys == [ids]
+    assert identifiable.id_type == "DATA_ROW_ID"
+
+
+def test__str__():
+    ids = ["a", "b", "c"]
+    identifiable = UniqueIds(ids)
+    str_representation = identifiable.__str__()
+    assert str_representation == "['a', 'b', 'c']"