improve result fetching experience

Author: Matt Sokoloff

labelbox/schema/task.py

@@ -1,7 +1,8 @@
 import logging
 import requests
 import time
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Optional, Dict, Any
+from functools import lru_cache
 
 from labelbox.exceptions import ResourceNotFoundError
 from labelbox.orm.db_object import DbObject
@@ -32,13 +33,14 @@ class Task(DbObject):
     name = Field.String("name")
     status = Field.String("status")
     completion_percentage = Field.Float("completion_percentage")
-    result = Field.String("result")
+    result_url = Field.String("result_url", "result")
     _user: Optional["User"] = None
 
     # Relationships
     created_by = Relationship.ToOne("User", False, "created_by")
     organization = Relationship.ToOne("Organization")
 
+
     def refresh(self) -> None:
         """ Refreshes Task data from the server. """
         assert self._user is not None
@@ -68,12 +70,36 @@ def wait_till_done(self, timeout_seconds=300) -> None:
             time.sleep(sleep_time_seconds)
             self.refresh()
 
-    def errors(self):
+    @property
+    def errors(self) -> Optional[Dict[str, Any]]:
         """ Downloads the result file from Task
         """
-        if self.status == "FAILED" and self.result:
-            response = requests.get(self.result)
-            response.raise_for_status()
-            data = response.json()
-            return data.get('error')
+        self.wait_till_done(timeout_seconds = 600)
+        if self.status == "FAILED":
+            data = self._fetch_remote(self.result_url)
+            if data:
+                return data.get('error', None)
+        elif self.status == "IN_PROGRESS":
+            raise Exception("Job state IN_PROGRESS. Result not available.")
         return None
+
+    @property
+    def result(self) -> Dict[str, Any]:
+        """ Fetch the result for a task
+        """
+        self.wait_till_done(timeout_seconds = 600)
+        if self.status == "COMPLETE":
+            return self._fetch_remote(self.result_url)
+        elif self.status == "FAILED":
+            raise Exception(f"Job failed. Errors : {self.errors()}")
+        else:
+            raise Exception("Job state IN_PROGRESS. Result not available.")
+
+    @lru_cache()
+    def _fetch_remote(self, result_url) -> Dict[str, Any]:
+        """ Function for fetching and caching the result data.
+        """
+        response = requests.get(result_url)
+        response.raise_for_status()
+        return response.json()
+

tests/integration/test_data_rows.py

@@ -543,3 +543,22 @@ def test_delete_data_row_attachment(datarow, image_url):
         attachment.delete()
 
     assert len(list(datarow.attachments())) == 0
+
+
+def test_create_data_rows_result(client, dataset, image_url):
+    task = dataset.create_data_rows([
+            {
+                DataRow.row_data: image_url,
+                DataRow.external_id: "row1",
+            },
+            {
+                DataRow.row_data: image_url,
+                DataRow.external_id: "row1",
+            },
+        ])
+    assert task.errors is None
+    assert 'createdDataRows' in task.result
+    results = task.result['createdDataRows']
+    for result in results:
+        client.get_data_row(result['id'])
+

tests/integration/test_task.py

@@ -22,8 +22,8 @@ def test_task_errors(dataset, image_url):
     assert task in client.get_user().created_tasks()
     task.wait_till_done()
     assert task.status == "FAILED"
-    assert task.errors() is not None
-    assert 'message' in task.errors()
+    assert task.errors is not None
+    assert 'message' in task.errors
 
 
 def test_task_success(dataset, image_url):
@@ -36,4 +36,6 @@ def test_task_success(dataset, image_url):
     assert task in client.get_user().created_tasks()
     task.wait_till_done()
     assert task.status == "COMPLETE"
-    assert task.errors() is None
+    assert task.errors is None
+    assert task.result is not None
+