Merge pull request #876 from Labelbox/mno/x-0-fix-default-values

Author: mnoszczak

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"]

labelbox/schema/project.py

@@ -393,18 +393,23 @@ 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,
             "data_row_details": False,
             "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} }

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."
             )