[AL-2075] Batch list and export

Author: farkob

labelbox/schema/batch.py

@@ -1,5 +1,14 @@
-from labelbox.orm.db_object import DbObject
-from labelbox.orm.model import Field, Relationship
+from typing import Generator
+from labelbox.orm.db_object import DbObject, experimental
+from labelbox.orm.model import Entity, Field, Relationship
+from labelbox.exceptions import LabelboxError
+from io import StringIO
+import ndjson
+import requests
+import logging
+import time
+
+logger = logging.getLogger(__name__)
 
 
 class Batch(DbObject):
@@ -19,7 +28,50 @@ class Batch(DbObject):
     created_at = Field.DateTime("created_at")
     updated_at = Field.DateTime("updated_at")
     size = Field.Int("size")
+    archived_at = Field.DateTime("archived_at")
 
     # Relationships
     project = Relationship.ToOne("Project")
     created_by = Relationship.ToOne("User")
+
+    def export_data_rows(self, timeout_seconds=120) -> Generator:
+        """ Returns a generator that produces all data rows that are currently
+        in this batch.
+
+        Note: For efficiency, the data are cached for 30 minutes. Newly created data rows will not appear
+        until the end of the cache period.
+
+        Args:
+            timeout_seconds (float): Max waiting time, in seconds.
+        Returns:
+            Generator that yields DataRow objects belonging to this batch.
+        Raises:
+            LabelboxError: if the export fails or is unable to download within the specified time.
+        """
+        id_param = "batchId"
+        query_str = """mutation GetBatchDataRowsExportUrlPyApi($%s: ID!)
+            {exportBatchDataRows(data:{batchId: $%s }) {downloadUrl createdAt status}}
+        """ % (id_param, id_param)
+        sleep_time = 2
+        while True:
+            res = self.client.execute(query_str, {id_param: self.uid})
+            res = res["exportBatchDataRows"]
+            if res["status"] == "COMPLETE":
+                download_url = res["downloadUrl"]
+                response = requests.get(download_url)
+                response.raise_for_status()
+                reader = ndjson.reader(StringIO(response.text))
+                return (
+                    Entity.DataRow(self.client, result) for result in reader)
+            elif res["status"] == "FAILED":
+                raise LabelboxError("Data row export failed.")
+
+            timeout_seconds -= sleep_time
+            if timeout_seconds <= 0:
+                raise LabelboxError(
+                    f"Unable to export data rows within {timeout_seconds} seconds."
+                )
+
+            logger.debug("Batch '%s' data row export, waiting for server...",
+                         self.uid)
+            time.sleep(sleep_time)

labelbox/schema/project.py

@@ -298,7 +298,7 @@ def export_labels(self,
 
         def _string_from_dict(dictionary: dict, value_with_quotes=False) -> str:
             """Returns a concatenated string of the dictionary's keys and values
-            
+
             The string will be formatted as {key}: 'value' for each key. Value will be inclusive of
             quotations while key will not. This can be toggled with `value_with_quotes`"""
 
@@ -840,6 +840,25 @@ def bulk_import_requests(self) -> PaginatedCollection:
                                    ["bulkImportRequests"],
                                    Entity.BulkImportRequest)
 
+    def batches(self) -> PaginatedCollection:
+        """ Fetch all batches that belong to this project
+
+        Returns:
+            A `PaginatedCollection of `Batch`es
+        """
+        id_param = "projectId"
+        query_str = """query GetProjectBatchesPyApi($from: String, $first: PageSize, $%s: ID!) {
+            project(where: {id: $%s}) {id
+            batches(after: $from, first: $first) { nodes { %s } pageInfo { endCursor }}}}
+        """ % (id_param, id_param, query.results_query_part(Entity.Batch))
+        return PaginatedCollection(self.client,
+                                   query_str, {id_param: self.uid},
+                                   ['project', 'batches', 'nodes'],
+                                   Entity.Batch,
+                                   cursor_path=['project', 'batches',
+                                                'pageInfo', 'endCursor'],
+                                   experimental=True)
+
     def upload_annotations(
             self,
             name: str,

tests/integration/test_batch.py

@@ -25,3 +25,22 @@ def test_create_batch(configured_project: Project, big_dataset: Dataset):
     batch = configured_project.create_batch("test-batch", data_rows, 3)
     assert batch.name == 'test-batch'
     assert batch.size == len(data_rows)
+
+def test_export_data_rows(configured_project: Project, dataset: Dataset):
+    n_data_rows = 5
+    task = dataset.create_data_rows([
+        {
+            "row_data": IMAGE_URL,
+            "external_id": "my-image"
+        },
+    ] * n_data_rows)
+    task.wait_till_done()
+
+    data_rows = [dr.uid for dr in list(big_dataset.export_data_rows())]
+    batch = configured_project.create_batch("batch test", data_rows)
+
+    result = list(batch.export_data_rows())
+    exported_data_rows = [dr.uid for dr in result]
+
+    assert len(result) == n_data_rows
+    assert set(data_rows) == set(exported_data_rows)