Merge pull request #777 from Labelbox/mkozik/create-batch-with-data-rows

[X-0] Fix create batch with data row classes

Author: mkozik-lb

labelbox/schema/project.py

@@ -21,6 +21,7 @@
 from labelbox.schema.media_type import MediaType
 from labelbox.schema.queue_mode import QueueMode
 from labelbox.schema.resource_tag import ResourceTag
+from labelbox.schema.data_row import DataRow
 
 if TYPE_CHECKING:
     from labelbox import BulkImportRequest
@@ -574,7 +575,7 @@ def setup(self, labeling_frontend, labeling_frontend_options) -> None:
 
     def create_batch(self,
                      name: str,
-                     data_rows: List[str],
+                     data_rows: List[Union[str, DataRow]],
                      priority: int = 5,
                      consensus_settings: Optional[Dict[str, float]] = None):
         """Create a new batch for a project. Batches is in Beta and subject to change
@@ -606,7 +607,7 @@ def create_batch(self,
             raise ValueError("You need at least one data row in a batch")
 
         self._wait_until_data_rows_are_processed(
-            data_rows, self._wait_processing_max_seconds)
+            dr_ids, self._wait_processing_max_seconds)
         method = 'createBatchV2'
         query_str = """mutation %sPyApi($projectId: ID!, $batchInput: CreateBatchInput!) {
               project(where: {id: $projectId}) {

tests/integration/test_batch.py

@@ -56,8 +56,8 @@ def test_create_batch(batch_project: Project, big_dataset: Dataset):
 
 
 def test_create_batch_with_consensus_settings(batch_project: Project,
-                                              big_dataset: Dataset):
-    data_rows = [dr.uid for dr in list(big_dataset.export_data_rows())]
+                                              small_dataset: Dataset):
+    data_rows = [dr.uid for dr in list(small_dataset.export_data_rows())]
     consensus_settings = {"coverage_percentage": 0.1, "number_of_labels": 3}
     batch = batch_project.create_batch("batch with consensus settings",
                                        data_rows,
@@ -68,6 +68,14 @@ def test_create_batch_with_consensus_settings(batch_project: Project,
     assert batch.consensus_settings == consensus_settings
 
 
+def test_create_batch_with_data_row_class(batch_project: Project,
+                                          small_dataset: Dataset):
+    data_rows = list(small_dataset.export_data_rows())
+    batch = batch_project.create_batch("test-batch-data-rows", data_rows, 3)
+    assert batch.name == "test-batch-data-rows"
+    assert batch.size == len(data_rows)
+
+
 def test_archive_batch(batch_project: Project, small_dataset: Dataset):
     data_rows = [dr.uid for dr in list(small_dataset.export_data_rows())]
     batch = batch_project.create_batch("batch to archive", data_rows)