diff --git a/labelbox/schema/model_run.py b/labelbox/schema/model_run.py index 701fd38b7..56a42c336 100644 --- a/labelbox/schema/model_run.py +++ b/labelbox/schema/model_run.py @@ -456,14 +456,19 @@ def export_labels( """ - def export_v2(self, task_name: str, - params: Optional[ModelRunExportParams]) -> Task: - _params = params or {} + def export_v2(self, + task_name: Optional[str] = None, + params: Optional[ModelRunExportParams] = None) -> Task: mutation_name = "exportDataRowsInModelRun" create_task_query_str = """mutation exportDataRowsInModelRunPyApi($input: ExportDataRowsInModelRunInput!){ %s(input: $input) {taskId} } """ % (mutation_name) - params = { + if (task_name is None): + task_name = f'Export Data Rows in Model Run - {self.name}' + + _params = params or ModelRunExportParams() + + queryParams = { "input": { "taskName": task_name, "filters": { @@ -490,7 +495,7 @@ def export_v2(self, task_name: str, } res = self.client.execute( create_task_query_str, - params, + queryParams, ) res = res[mutation_name] task_id = res["taskId"] diff --git a/labelbox/schema/project.py b/labelbox/schema/project.py index 277bf6622..a7f596001 100644 --- a/labelbox/schema/project.py +++ b/labelbox/schema/project.py @@ -382,9 +382,14 @@ def _validate_datetime(string_date: str) -> bool: """ - def export_v2(self, task_name: str, - params: Optional[ProjectExportParams]) -> Task: - defaultParams: ProjectExportParams = { + def export_v2(self, + task_name: Optional[str] = None, + params: Optional[ProjectExportParams] = None) -> Task: + + if (task_name is None): + task_name = f'Export Data Rows in Project - {self.name}' + + _params = params or ProjectExportParams({ "attachments": False, "media_attributes": False, "metadata_fields": False, @@ -392,8 +397,8 @@ def export_v2(self, task_name: str, "project_details": False, "labels": False, "performance_details": False - } - _params: ProjectExportParams = params if params is not None else defaultParams + }) + mutation_name = "exportDataRowsInProject" create_task_query_str = """mutation exportDataRowsInProjectPyApi($input: ExportDataRowsInProjectInput!){ %s(input: $input) {taskId} } diff --git a/labelbox/schema/task.py b/labelbox/schema/task.py index 7ce4279e2..31f3390ee 100644 --- a/labelbox/schema/task.py +++ b/labelbox/schema/task.py @@ -1,7 +1,8 @@ import logging import requests import time -from typing import TYPE_CHECKING, Callable, Optional, Dict, Any, List +from typing import TYPE_CHECKING, Callable, Optional, Dict, Any, List, Union +import ndjson from labelbox.exceptions import ResourceNotFoundError from labelbox.orm.db_object import DbObject @@ -38,6 +39,7 @@ class Task(DbObject): status = Field.String("status") completion_percentage = Field.Float("completion_percentage") result_url = Field.String("result_url", "result") + type = Field.String("type") _user: Optional["User"] = None # Relationships @@ -92,13 +94,16 @@ def errors(self) -> Optional[Dict[str, Any]]: return None @property - def result(self) -> List[Dict[str, Any]]: + def result(self) -> Union[List[Dict[str, Any]], Dict[str, Any]]: """ Fetch the result for an import task. """ if self.status == "FAILED": raise ValueError(f"Job failed. Errors : {self.errors}") else: result = self._fetch_remote_json() + if self.type == 'export-data-rows': + return result + return [{ 'id': data_row['id'], 'external_id': data_row.get('externalId'), @@ -124,11 +129,18 @@ def _fetch_remote_json(self) -> Dict[str, Any]: def download_result(): response = requests.get(self.result_url) response.raise_for_status() - return response.json() - - if self.name != 'JSON Import': + try: + return response.json() + except Exception as e: + pass + try: + return ndjson.loads(response.text) + except Exception as e: + raise ValueError("Failed to parse task JSON/NDJSON result.") + + if self.name != 'JSON Import' and self.type != 'export-data-rows': raise ValueError( - "Task result is only supported for `JSON Import` tasks." + "Task result is only supported for `JSON Import` and `export` tasks." " Download task.result_url manually to access the result for other tasks." )