From f4ca4572b6fa6b7ee05dfd91241dcbeba844d8fa Mon Sep 17 00:00:00 2001
From: paulnoirel <87332996+paulnoirel@users.noreply.github.com>
Date: Fri, 14 Mar 2025 21:36:28 +0000
Subject: [PATCH 1/4] Add tasks to organization and TaskStatus to filter tasks

---
 libs/labelbox/src/labelbox/__init__.py        |  1 +
 libs/labelbox/src/labelbox/schema/__init__.py |  1 +
 .../src/labelbox/schema/organization.py       |  1 +
 libs/labelbox/src/labelbox/schema/task.py     |  4 +++
 .../src/labelbox/schema/taskstatus.py         | 25 +++++++++++++++
 .../test_export_data_rows_streamable.py       | 32 ++++++++++++++++++-
 6 files changed, 63 insertions(+), 1 deletion(-)
 create mode 100644 libs/labelbox/src/labelbox/schema/taskstatus.py

diff --git a/libs/labelbox/src/labelbox/__init__.py b/libs/labelbox/src/labelbox/__init__.py
index 663028916..42366d049 100644
--- a/libs/labelbox/src/labelbox/__init__.py
+++ b/libs/labelbox/src/labelbox/__init__.py
@@ -98,3 +98,4 @@
     PromptResponseClassification,
 )
 from lbox.exceptions import *
+from labelbox.schema.taskstatus import TaskStatus
diff --git a/libs/labelbox/src/labelbox/schema/__init__.py b/libs/labelbox/src/labelbox/schema/__init__.py
index d6b74de68..b22a54854 100644
--- a/libs/labelbox/src/labelbox/schema/__init__.py
+++ b/libs/labelbox/src/labelbox/schema/__init__.py
@@ -26,3 +26,4 @@
 import labelbox.schema.catalog
 import labelbox.schema.ontology_kind
 import labelbox.schema.project_overview
+import labelbox.schema.taskstatus
diff --git a/libs/labelbox/src/labelbox/schema/organization.py b/libs/labelbox/src/labelbox/schema/organization.py
index bd416e997..1eea3aebf 100644
--- a/libs/labelbox/src/labelbox/schema/organization.py
+++ b/libs/labelbox/src/labelbox/schema/organization.py
@@ -52,6 +52,7 @@ def __init__(self, *args, **kwargs):
     projects = Relationship.ToMany("Project", True)
     webhooks = Relationship.ToMany("Webhook", False)
     resource_tags = Relationship.ToMany("ResourceTags", False)
+    tasks = Relationship.ToMany("Task", False, "tasks")
 
     def invite_user(
         self,
diff --git a/libs/labelbox/src/labelbox/schema/task.py b/libs/labelbox/src/labelbox/schema/task.py
index d536b2560..6a58f99e8 100644
--- a/libs/labelbox/src/labelbox/schema/task.py
+++ b/libs/labelbox/src/labelbox/schema/task.py
@@ -14,6 +14,7 @@
 from labelbox.schema.internal.datarow_upload_constants import (
     DOWNLOAD_RESULT_PAGE_SIZE,
 )
+from labelbox.schema.taskstatus import TaskStatus
 
 if TYPE_CHECKING:
     from labelbox import User
@@ -45,6 +46,9 @@ class Task(DbObject):
     created_at = Field.DateTime("created_at")
     name = Field.String("name")
     status = Field.String("status")
+    status_type = Field.Enum(
+        TaskStatus, "status_type", "status"
+    )  # additional status for filtering
     completion_percentage = Field.Float("completion_percentage")
     result_url = Field.String("result_url", "result")
     errors_url = Field.String("errors_url", "errors")
diff --git a/libs/labelbox/src/labelbox/schema/taskstatus.py b/libs/labelbox/src/labelbox/schema/taskstatus.py
new file mode 100644
index 000000000..0abbce1ca
--- /dev/null
+++ b/libs/labelbox/src/labelbox/schema/taskstatus.py
@@ -0,0 +1,25 @@
+from enum import Enum
+
+
+class TaskStatus(str, Enum):
+    In_Progress = "IN_PROGRESS"
+    Complete = "COMPLETE"
+    Canceling = "CANCELLING"
+    Canceled = "CANCELED"
+    Failed = "FAILED"
+    Unknown = "UNKNOWN"
+
+    @classmethod
+    def _missing_(cls, value):
+        """Handle missing or unknown task status values.
+
+        If a task status value is not found in the enum, this method returns
+        the Unknown status instead of raising an error.
+
+        Args:
+            value: The status value that doesn't match any enum member
+
+        Returns:
+            TaskStatus.Unknown: The default status for unrecognized values
+        """
+        return cls.Unknown
diff --git a/libs/labelbox/tests/data/export/streamable/test_export_data_rows_streamable.py b/libs/labelbox/tests/data/export/streamable/test_export_data_rows_streamable.py
index 5ab4d30ec..f1a4e3840 100644
--- a/libs/labelbox/tests/data/export/streamable/test_export_data_rows_streamable.py
+++ b/libs/labelbox/tests/data/export/streamable/test_export_data_rows_streamable.py
@@ -1,6 +1,6 @@
 import time
 
-from labelbox import DataRow, ExportTask, StreamType
+from labelbox import DataRow, ExportTask, StreamType, Task, TaskStatus
 
 
 class TestExportDataRow:
@@ -135,3 +135,33 @@ def test_cancel_export_task(
         # Verify the task was cancelled
         cancelled_task = client.get_task_by_id(export_task.uid)
         assert cancelled_task.status in ["CANCELING", "CANCELED"]
+
+    def test_task_filter(self, client, data_row, wait_for_data_row_processing):
+        organization = client.get_organization()
+        user = client.get_user()
+
+        export_task = DataRow.export(
+            client=client,
+            data_rows=[data_row],
+            task_name="TestExportDataRow:test_task_filter",
+        )
+
+        # Check if task is listed "in progress" in organization's tasks
+        org_tasks_in_progress = organization.tasks(
+            where=Task.status_type == TaskStatus.In_Progress
+        )
+        retrieved_task_in_progress = next(
+            (t for t in org_tasks_in_progress if t.uid == export_task.uid), ""
+        )
+        assert getattr(retrieved_task_in_progress, "uid", "") == export_task.uid
+
+        export_task.wait_till_done()
+
+        # Check if task is listed "complete" in user's created tasks
+        user_tasks_complete = user.created_tasks(
+            where=Task.status_type == TaskStatus.Complete
+        )
+        retrieved_task_complete = next(
+            (t for t in user_tasks_complete if t.uid == export_task.uid), ""
+        )
+        assert getattr(retrieved_task_complete, "uid", "") == export_task.uid

From 6a79d460b96353d9be7270f96780143ac95cbca4 Mon Sep 17 00:00:00 2001
From: paulnoirel <87332996+paulnoirel@users.noreply.github.com>
Date: Thu, 20 Mar 2025 12:22:15 +0000
Subject: [PATCH 2/4] Add cancelable types and additional check to
 cancel_task()

---
 libs/labelbox/src/labelbox/client.py | 48 ++++++++++++++++++++++++++--
 1 file changed, 46 insertions(+), 2 deletions(-)

diff --git a/libs/labelbox/src/labelbox/client.py b/libs/labelbox/src/labelbox/client.py
index 956f12487..8c63d559a 100644
--- a/libs/labelbox/src/labelbox/client.py
+++ b/libs/labelbox/src/labelbox/client.py
@@ -79,6 +79,7 @@
 from labelbox.schema.slice import CatalogSlice, ModelSlice
 from labelbox.schema.task import DataUpsertTask, Task
 from labelbox.schema.user import User
+from labelbox.schema.taskstatus import TaskStatus
 
 logger = logging.getLogger(__name__)
 
@@ -90,6 +91,9 @@ class Client:
     top-level data objects (Projects, Datasets).
     """
 
+    # Class variable to cache task types
+    _cancelable_task_types = None
+
     def __init__(
         self,
         api_key=None,
@@ -2390,9 +2394,31 @@ def get_task_by_id(self, task_id: str) -> Union[Task, DataUpsertTask]:
         task._user = user
         return task
 
+    def _get_cancelable_task_types(self):
+        """Internal method that returns a list of task types that can be canceled.
+
+        The result is cached after the first call to avoid unnecessary API requests.
+
+        Returns:
+            List[str]: List of cancelable task types in snake_case format
+        """
+        if self._cancelable_task_types is None:
+            query = """query GetCancelableTaskTypes {
+                cancelableTaskTypes
+            }"""
+
+            result = self.execute(query).get("cancelableTaskTypes", [])
+            # Reformat to kebab case
+            self._cancelable_task_types = [
+                utils.snake_case(task_type).replace("_", "-")
+                for task_type in result
+            ]
+
+        return self._cancelable_task_types
+
     def cancel_task(self, task_id: str) -> bool:
         """
-        Cancels a task with the given ID.
+        Cancels a task with the given ID if the task type is cancelable and the task is in progress.
 
         Args:
             task_id (str): The ID of the task to cancel.
@@ -2401,8 +2427,26 @@ def cancel_task(self, task_id: str) -> bool:
             bool: True if the task was successfully cancelled.
 
         Raises:
-            LabelboxError: If the task could not be cancelled.
+            LabelboxError: If the task could not be cancelled, if the task type is not cancelable,
+                or if the task is not in progress.
+            ResourceNotFoundError: If the task does not exist (raised by get_task_by_id).
         """
+        # Get the task object to check its type and status
+        task = self.get_task_by_id(task_id)
+
+        # Check if task type is cancelable
+        cancelable_types = self._get_cancelable_task_types()
+        if task.type not in cancelable_types:
+            raise LabelboxError(
+                f"Task type '{task.type}' cannot be cancelled. Cancelable types are: {cancelable_types}"
+            )
+
+        # Check if task is in progress
+        if task.status_type != TaskStatus.In_Progress:
+            raise LabelboxError(
+                f"Task cannot be cancelled because it is not in progress. Current status: {task.status}"
+            )
+
         mutation_str = """
         mutation CancelTaskPyApi($id: ID!) {
             cancelBulkOperationJob(id: $id) {

From 8926ef461b205fe3cecbbd1469befaa037edac6b Mon Sep 17 00:00:00 2001
From: paulnoirel <87332996+paulnoirel@users.noreply.github.com>
Date: Thu, 20 Mar 2025 13:04:31 +0000
Subject: [PATCH 3/4] Refactoring: replace Task.status_type with
 Task.status_as_enum

---
 libs/labelbox/src/labelbox/client.py                          | 2 +-
 libs/labelbox/src/labelbox/schema/task.py                     | 4 ++--
 .../export/streamable/test_export_data_rows_streamable.py     | 4 ++--
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/libs/labelbox/src/labelbox/client.py b/libs/labelbox/src/labelbox/client.py
index 8c63d559a..6a88d790a 100644
--- a/libs/labelbox/src/labelbox/client.py
+++ b/libs/labelbox/src/labelbox/client.py
@@ -2442,7 +2442,7 @@ def cancel_task(self, task_id: str) -> bool:
             )
 
         # Check if task is in progress
-        if task.status_type != TaskStatus.In_Progress:
+        if task.status_as_enum != TaskStatus.In_Progress:
             raise LabelboxError(
                 f"Task cannot be cancelled because it is not in progress. Current status: {task.status}"
             )
diff --git a/libs/labelbox/src/labelbox/schema/task.py b/libs/labelbox/src/labelbox/schema/task.py
index 6a58f99e8..04fd7b12f 100644
--- a/libs/labelbox/src/labelbox/schema/task.py
+++ b/libs/labelbox/src/labelbox/schema/task.py
@@ -46,8 +46,8 @@ class Task(DbObject):
     created_at = Field.DateTime("created_at")
     name = Field.String("name")
     status = Field.String("status")
-    status_type = Field.Enum(
-        TaskStatus, "status_type", "status"
+    status_as_enum = Field.Enum(
+        TaskStatus, "status_as_enum", "status"
     )  # additional status for filtering
     completion_percentage = Field.Float("completion_percentage")
     result_url = Field.String("result_url", "result")
diff --git a/libs/labelbox/tests/data/export/streamable/test_export_data_rows_streamable.py b/libs/labelbox/tests/data/export/streamable/test_export_data_rows_streamable.py
index f1a4e3840..233fc2144 100644
--- a/libs/labelbox/tests/data/export/streamable/test_export_data_rows_streamable.py
+++ b/libs/labelbox/tests/data/export/streamable/test_export_data_rows_streamable.py
@@ -148,7 +148,7 @@ def test_task_filter(self, client, data_row, wait_for_data_row_processing):
 
         # Check if task is listed "in progress" in organization's tasks
         org_tasks_in_progress = organization.tasks(
-            where=Task.status_type == TaskStatus.In_Progress
+            where=Task.status_as_enum == TaskStatus.In_Progress
         )
         retrieved_task_in_progress = next(
             (t for t in org_tasks_in_progress if t.uid == export_task.uid), ""
@@ -159,7 +159,7 @@ def test_task_filter(self, client, data_row, wait_for_data_row_processing):
 
         # Check if task is listed "complete" in user's created tasks
         user_tasks_complete = user.created_tasks(
-            where=Task.status_type == TaskStatus.Complete
+            where=Task.status_as_enum == TaskStatus.Complete
         )
         retrieved_task_complete = next(
             (t for t in user_tasks_complete if t.uid == export_task.uid), ""

From e4e2ab39b20a57964438c97e3cd55906642292a6 Mon Sep 17 00:00:00 2001
From: paulnoirel <87332996+paulnoirel@users.noreply.github.com>
Date: Thu, 20 Mar 2025 14:09:29 +0000
Subject: [PATCH 4/4] Fix GraphQL query name: GetCancelableTaskTypesPyApi

---
 libs/labelbox/src/labelbox/client.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/libs/labelbox/src/labelbox/client.py b/libs/labelbox/src/labelbox/client.py
index 6a88d790a..1926957a5 100644
--- a/libs/labelbox/src/labelbox/client.py
+++ b/libs/labelbox/src/labelbox/client.py
@@ -2403,7 +2403,7 @@ def _get_cancelable_task_types(self):
             List[str]: List of cancelable task types in snake_case format
         """
         if self._cancelable_task_types is None:
-            query = """query GetCancelableTaskTypes {
+            query = """query GetCancelableTaskTypesPyApi {
                 cancelableTaskTypes
             }"""