Merge pull request #156 from Labelbox/ms/download-export

download exported data rows

Author: msokoloff1

labelbox/schema/project.py

@@ -6,6 +6,8 @@
 from pathlib import Path
 from typing import Dict, Union, Iterable
 from urllib.parse import urlparse
+import requests
+import ndjson
 
 from labelbox import utils
 from labelbox.schema.data_row import DataRow
@@ -166,8 +168,9 @@ def export_queued_data_rows(self, timeout_seconds=120):
         Args:
             timeout_seconds (float): Max waiting time, in seconds.
         Returns:
-            URL of the data file with this DataRow information. If the server didn't
-            generate during the `timeout_seconds` period, None is returned.
+            Data row fields for all data rows in the queue as json
+        Raises:
+            LabelboxError: if the export fails or is unable to download within the specified time.
         """
         id_param = "projectId"
         query_str = """mutation GetQueuedDataRowsExportUrlPyApi($%s: ID!)
@@ -178,7 +181,10 @@ def export_queued_data_rows(self, timeout_seconds=120):
             res = self.client.execute(query_str, {id_param: self.uid})
             res = res["exportQueuedDataRows"]
             if res["status"] == "COMPLETE":
-                return res["downloadUrl"]
+                download_url = res["downloadUrl"]
+                response = requests.get(download_url)
+                response.raise_for_status()
+                return ndjson.loads(response.text)
             elif res["status"] == "FAILED":
                 raise LabelboxError("Data row export failed.")
 

tests/integration/test_project.py

@@ -105,7 +105,5 @@ def test_attach_instructions(client, project):
 
 
 def test_queued_data_row_export(configured_project):
-    url = configured_project.export_queued_data_rows()
-    assert url
-    result = ndjson.loads(requests.get(url).text)
+    result = configured_project.export_queued_data_rows()
     assert len(result) == 1