Add DisconncetedTask

Author: Val Brodsky

docs/labelbox/disconnected-task.rst

@@ -0,0 +1,7 @@
+Disconnected Task
+===============================================================================================
+
+.. automodule:: labelbox.schema.disconnected_task
+   :members:
+   :exclude-members: SupportsTaskQueries, TaskFactory
+   :show-inheritance:

libs/labelbox/src/labelbox/schema/disconnected_task.py

@@ -0,0 +1,93 @@
+from typing import Any, Dict, List, Optional, Protocol, Union
+
+from labelbox.exceptions import OperationNotSupportedException, ResourceNotFoundError
+from labelbox.orm.db_object import experimental
+from labelbox.orm.model import Field
+from labelbox.schema.task import Task, DataUpsertTask
+
+TaskDataType = Optional[Union[Dict[str, Any], List[Dict[str, Any]]]]
+
+
+class SupportsTaskQueries(Protocol):
+
+    @property
+    def status(self) -> Field:
+        ...
+
+    @property
+    def errors(self) -> TaskDataType:
+        ...
+
+    @property
+    def result(self) -> TaskDataType:
+        ...
+
+
+class TaskFactory:
+
+    @classmethod
+    def get(cls, client, task_id) -> SupportsTaskQueries:
+        return cls.get_task_by_id(client, task_id)
+
+    @classmethod
+    def get_task_by_id(cls, client, task_id) -> SupportsTaskQueries:
+        user = client.get_user()
+        query = """
+            query GetUserCreatedTasksPyApi($userId: ID!, $taskId: ID!) {
+            user(where: {id: $userId}) {
+                createdTasks(where: {id: $taskId} skip: 0 first: 1) {
+                    completionPercentage
+                    createdAt
+                    errors
+                    metadata
+                    name
+                    result
+                    status
+                    type
+                    id
+                    updatedAt
+                }
+            }
+        }
+        """
+        result = client.execute(query, {"userId": user.uid, "taskId": task_id})
+        data = result.get("user", {}).get("createdTasks", [])
+        if not data:
+            raise ResourceNotFoundError(
+                message=f"The task {task_id} does not exist.")
+        task_data = data[0]
+        if task_data["type"].lower() == 'adv-upsert-data-rows':
+            task = DataUpsertTask(client, task_data)
+        else:
+            task = Task(client, task_data)
+
+        task._user = user
+        return task
+
+
+class DisconnectedTask:
+    """
+    A class to interact with a task by task id.
+
+    This class can be used to avoid waiting to tasks that take too long
+    """
+
+    @experimental
+    def __init__(self, client, task_id: str):
+        self.task = TaskFactory.get(client, task_id)
+
+    def is_completed(self) -> bool:
+        return str(self.task.status) != "IN_PROGRESS"
+
+    def status(self) -> str:
+        return str(self.task.status)
+
+    def errors(self) -> TaskDataType:
+        if not self.is_completed():
+            raise OperationNotSupportedException("Task is not completed yet")
+        return self.task.errors
+
+    def result(self) -> TaskDataType:
+        if not self.is_completed():
+            raise OperationNotSupportedException("Task is not completed yet")
+        return self.task.result

libs/labelbox/src/labelbox/schema/task.py

@@ -2,7 +2,7 @@
 import logging
 import requests
 import time
-from typing import TYPE_CHECKING, Callable, Optional, Dict, Any, List, Union, Final
+from typing import TYPE_CHECKING, Callable, Optional, Dict, Any, List, Union
 from labelbox import parser
 
 from labelbox.exceptions import ResourceNotFoundError

libs/labelbox/tests/integration/test_task.py

@@ -1,4 +1,5 @@
 import json
+from labelbox.schema.disconnected_task import DisconnectedTask
 import pytest
 import collections.abc
 from labelbox import DataRow
@@ -31,6 +32,14 @@ def test_task_errors(dataset, image_url, snapshot):
         0]['message']
     assert len(task.failed_data_rows[0]['failedDataRows'][0]['metadata']) == 2
 
+    dt = DisconnectedTask(client, task.uid)
+    assert dt.is_completed()
+    assert dt.status() == "COMPLETE"
+    assert len(dt.errors()) == 1
+    assert dt.errors()[0]['message'].startswith(
+        "A schemaId can only be specified once per DataRow")
+    assert dt.result() is None
+
 
 def test_task_success_json(dataset, image_url, snapshot):
     client = dataset.client
@@ -57,3 +66,9 @@ def test_task_success_json(dataset, image_url, snapshot):
     snapshot.assert_match(json.dumps(task_result),
                           'test_task.test_task_success_json.json')
     assert len(task.result)
+
+    dt = DisconnectedTask(client, task.uid)
+    assert dt.is_completed()
+    assert dt.status() == "COMPLETE"
+    assert len(dt.result()) == 1
+    assert dt.errors() is None