Merge pull request #593 from Labelbox/ms/fetch-result

improve result fetching experience

Author: msokoloff1

labelbox/schema/project.py

@@ -327,7 +327,7 @@ def _validate_datetime(string_date: str) -> bool:
                 try:
                     datetime.strptime(string_date, "%Y-%m-%d")
                 except ValueError:
-                    raise ValueError(f"""Incorrect format for: {string_date}. 
+                    raise ValueError(f"""Incorrect format for: {string_date}.
                     Format must be \"YYYY-MM-DD\"""")
             return True
 

labelbox/schema/task.py

@@ -1,7 +1,8 @@
 import logging
+import json
 import requests
 import time
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, TypeVar, Callable, Optional, Dict, Any, List
 
 from labelbox.exceptions import ResourceNotFoundError
 from labelbox.orm.db_object import DbObject
@@ -10,6 +11,11 @@
 if TYPE_CHECKING:
     from labelbox import User
 
+    def lru_cache() -> Callable[..., Callable[..., Dict[str, Any]]]:
+        pass
+else:
+    from functools import lru_cache
+
 logger = logging.getLogger(__name__)
 
 
@@ -32,7 +38,7 @@ 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
@@ -68,12 +74,44 @@ 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')
+        if self.status == "FAILED":
+            result = self._fetch_remote_json()
+            return result['error']
         return None
+
+    @property
+    def result(self) -> List[Dict[str, Any]]:
+        """ Fetch the result for a task
+        """
+        if self.status == "FAILED":
+            raise ValueError(f"Job failed. Errors : {self.errors}")
+        else:
+            result = self._fetch_remote_json()
+            return [{
+                'id': data_row['id'],
+                'external_id': data_row.get('externalId'),
+                'row_data': data_row['rowData']
+            } for data_row in result['createdDataRows']]
+
+    @lru_cache()
+    def _fetch_remote_json(self) -> Dict[str, Any]:
+        """ Function for fetching and caching the result data.
+        """
+        if self.name != 'JSON Import':
+            raise ValueError(
+                "Task result is only supported for `JSON Import` tasks."
+                " Download task.result_url manually to access the result for other tasks."
+            )
+        self.wait_till_done(timeout_seconds=600)
+        if self.status == "IN_PROGRESS":
+            raise ValueError(
+                "Job status still in `IN_PROGRESS`. The result is not available. Call task.wait_till_done() with a larger timeout or contact support."
+            )
+
+        response = requests.get(self.result_url)
+        response.raise_for_status()
+        return response.json()

tests/integration/test_data_rows.py

@@ -543,3 +543,19 @@ 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
+    for result in task.result:
+        client.get_data_row(result['id'])

tests/integration/test_task.py

@@ -1,4 +1,6 @@
-from labelbox import DataRow
+import pytest
+
+from labelbox import DataRow, Task
 from labelbox.schema.data_row_metadata import DataRowMetadataField
 
 EMBEDDING_SCHEMA_ID = "ckpyije740000yxdk81pbgjdc"
@@ -22,11 +24,14 @@ 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
+    with pytest.raises(Exception) as exc_info:
+        task.result
+    assert str(exc_info.value).startswith("Job failed. Errors : {")
 
 
-def test_task_success(dataset, image_url):
+def test_task_success_json(dataset, image_url):
     client = dataset.client
     task = dataset.create_data_rows([
         {
@@ -36,4 +41,20 @@ 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
+    assert len(task.result)
+
+
+def test_task_success_label_export(client, configured_project_with_label):
+    project, _, _, _ = configured_project_with_label
+    project.export_labels()
+    user = client.get_user()
+    task = None
+    for task in user.created_tasks():
+        if task.name != 'JSON Import':
+            break
+
+    with pytest.raises(ValueError) as exc_info:
+        task.result
+    assert str(exc_info.value).startswith("Task result is only supported for")