PR feedback: refactor DataRowUpsertItem

Author: Val Brodsky

libs/labelbox/src/labelbox/schema/dataset.py

@@ -32,9 +32,7 @@
 from labelbox.schema.task import Task, DataUpsertTask
 from labelbox.schema.user import User
 from labelbox.schema.iam_integration import IAMIntegration
-from labelbox.schema.internal.data_row_create_upsert import (DataRowItemBase,
-                                                             DataRowUpsertItem,
-                                                             DataRowCreateItem)
+from labelbox.schema.internal.data_row_upsert_item import (DataRowUpsertItem)
 from labelbox.schema.internal.data_row_uploader import DataRowUploader
 from labelbox.schema.internal.datarow_upload_constants import (
     MAX_DATAROW_PER_API_OPERATION, FILE_UPLOAD_THREAD_COUNT, UPSERT_CHUNK_SIZE)
@@ -290,7 +288,7 @@ def create_data_rows(self,
 
         if len(string_items) > 0:
             dict_string_items = self._build_from_local_paths(string_items)
-        specs = DataRowCreateItem.build(self.uid,
+        specs = DataRowUpsertItem.build(self.uid,
                                         dict_items + dict_string_items)
         return self._exec_upsert_data_rows(specs, file_upload_thread_count)
 
@@ -607,12 +605,12 @@ def upsert_data_rows(self,
         >>>     ])
         >>>     task.wait_till_done()
         """
-        specs = DataRowUpsertItem.build(self.uid, items)
+        specs = DataRowUpsertItem.build(self.uid, items, (UniqueId, GlobalKey))
         return self._exec_upsert_data_rows(specs, file_upload_thread_count)
 
     def _exec_upsert_data_rows(
         self,
-        specs: List[DataRowItemBase],
+        specs: List[DataRowUpsertItem],
         file_upload_thread_count: int = FILE_UPLOAD_THREAD_COUNT
     ) -> "DataUpsertTask":
 

libs/labelbox/src/labelbox/schema/internal/data_row_uploader.py

@@ -10,8 +10,8 @@
 from labelbox.orm.model import Entity
 from labelbox.orm.model import Field
 from labelbox.schema.embedding import EmbeddingVector
-from labelbox.schema.internal.data_row_create_upsert import DataRowItemBase
 from labelbox.schema.internal.datarow_upload_constants import MAX_DATAROW_PER_API_OPERATION
+from labelbox.schema.internal.data_row_upsert_item import DataRowUpsertItem
 
 
 class UploadManifest:
@@ -213,7 +213,7 @@ def formatLegacyConversationalData(item):
             return item
 
         def convert_item(data_row_item):
-            if isinstance(data_row_item, DataRowItemBase):
+            if isinstance(data_row_item, DataRowUpsertItem):
                 item = data_row_item.payload
             else:
                 item = data_row_item
@@ -238,7 +238,7 @@ def convert_item(data_row_item):
             # Upload any local file paths
             item = upload_if_necessary(item)
 
-            if isinstance(data_row_item, DataRowItemBase):
+            if isinstance(data_row_item, DataRowUpsertItem):
                 return {'id': data_row_item.id, 'payload': item}
             else:
                 return item
@@ -263,7 +263,7 @@ def convert_item(data_row_item):
                                   filename="json_import.json")
 
     @staticmethod
-    def upload_in_chunks(client, specs: List[DataRowItemBase],
+    def upload_in_chunks(client, specs: List[DataRowUpsertItem],
                          file_upload_thread_count: int,
                          upsert_chunk_size: int) -> UploadManifest:
         empty_specs = list(filter(lambda spec: spec.is_empty(), specs))

libs/labelbox/src/labelbox/schema/internal/data_row_create_upsert.py renamed to libs/labelbox/src/labelbox/schema/internal/data_row_upsert_item.py

@@ -1,19 +1,17 @@
-from abc import ABC, abstractmethod
 from typing import List, Tuple, Optional
 
 from labelbox.schema.identifiable import UniqueId, GlobalKey
 from labelbox.pydantic_compat import BaseModel
 
 
-class DataRowItemBase(BaseModel, ABC):
+class DataRowUpsertItem(BaseModel):
     """
     Base class for creating payloads for upsert operations.
     """
     id: dict
     payload: dict
 
     @classmethod
-    @abstractmethod
     def build(
         cls,
         dataset_id: str,
@@ -52,25 +50,24 @@ def is_empty(self) -> bool:
                 len(self.payload.keys()) == 1 and "dataset_id" in self.payload)
 
 
-class DataRowUpsertItem(DataRowItemBase):
-
-    @classmethod
-    def build(
-        cls,
-        dataset_id: str,
-        items: List[dict],
-        key_types: Optional[Tuple[type, ...]] = ()
-    ) -> List["DataRowItemBase"]:
-        return super().build(dataset_id, items, (UniqueId, GlobalKey))
+# class DataRowUpsertItem(DataRowItemBase):
 
+#     @classmethod
+#     def build(
+#         cls,
+#         dataset_id: str,
+#         items: List[dict],
+#         key_types: Optional[Tuple[type, ...]] = ()
+#     ) -> List["DataRowItemBase"]:
+#         return super().build(dataset_id, items, (UniqueId, GlobalKey))
 
-class DataRowCreateItem(DataRowItemBase):
+# class DataRowCreateItem(DataRowItemBase):
 
-    @classmethod
-    def build(
-        cls,
-        dataset_id: str,
-        items: List[dict],
-        key_types: Optional[Tuple[type, ...]] = ()
-    ) -> List["DataRowItemBase"]:
-        return super().build(dataset_id, items, ())
+#     @classmethod
+#     def build(
+#         cls,
+#         dataset_id: str,
+#         items: List[dict],
+#         key_types: Optional[Tuple[type, ...]] = ()
+#     ) -> List["DataRowItemBase"]:
+#         return super().build(dataset_id, items, ())

libs/labelbox/tests/unit/test_data_row_upsert_data.py

@@ -1,6 +1,5 @@
 import pytest
-from labelbox.schema.internal.data_row_create_upsert import (DataRowUpsertItem,
-                                                             DataRowCreateItem)
+from labelbox.schema.internal.data_row_upsert_item import (DataRowUpsertItem)
 from labelbox.schema.identifiable import UniqueId, GlobalKey
 from labelbox.schema.asset_attachment import AttachmentType
 
@@ -13,12 +12,15 @@ def data_row_create_items():
             "row_data": "http://my_site.com/photos/img_01.jpg",
             "global_key": "global_key1",
             "external_id": "ex_id1",
-            "attachments": [
-                {"type": AttachmentType.RAW_TEXT, "name": "att1", "value": "test1"}
-            ],
-            "metadata": [
-                {"name": "tag", "value": "tag value"},
-            ]
+            "attachments": [{
+                "type": AttachmentType.RAW_TEXT,
+                "name": "att1",
+                "value": "test1"
+            }],
+            "metadata": [{
+                "name": "tag",
+                "value": "tag value"
+            },]
         },
     ]
     return dataset_id, items
@@ -44,15 +46,15 @@ def test_data_row_upsert_items(data_row_create_items, data_row_update_items):
     dataset_id, create_items = data_row_create_items
     dataset_id, update_items = data_row_update_items
     items = create_items + update_items
-    result = DataRowUpsertItem.build(dataset_id, items)
+    result = DataRowUpsertItem.build(dataset_id, items, (UniqueId, GlobalKey))
     assert len(result) == len(items)
     for item, res in zip(items, result):
         assert res.payload == item
 
 
 def test_data_row_create_items(data_row_create_items):
     dataset_id, items = data_row_create_items
-    result = DataRowCreateItem.build(dataset_id, items)
+    result = DataRowUpsertItem.build(dataset_id, items)
     assert len(result) == len(items)
     for item, res in zip(items, result):
         assert res.payload == item
@@ -61,4 +63,4 @@ def test_data_row_create_items(data_row_create_items):
 def test_data_row_create_items_not_updateable(data_row_update_items):
     dataset_id, items = data_row_update_items
     with pytest.raises(ValueError):
-        DataRowCreateItem.build(dataset_id, items)
+        DataRowUpsertItem.build(dataset_id, items, ())