Merge pull request #841 from Labelbox/mno/AL-4397

Author: mnoszczak

labelbox/schema/export_params.py

@@ -16,7 +16,7 @@ class DataRowParams(TypedDict):
 
 class ProjectExportParams(DataRowParams):
     project_details: Optional[bool]
-    label_details: Optional[bool]
+    labels: Optional[bool]
     performance_details: Optional[bool]
 
 

labelbox/schema/project.py

@@ -12,16 +12,20 @@
 
 from labelbox import utils
 from labelbox.exceptions import (InvalidQueryError, LabelboxError,
-                                 ProcessingWaitTimeout, ResourceConflict)
+                                 ProcessingWaitTimeout, ResourceConflict,
+                                 ResourceNotFoundError)
 from labelbox.orm import query
 from labelbox.orm.db_object import DbObject, Deletable, Updateable
 from labelbox.orm.model import Entity, Field, Relationship
 from labelbox.pagination import PaginatedCollection
 from labelbox.schema.consensus_settings import ConsensusSettings
 from labelbox.schema.data_row import DataRow
+from labelbox.schema.export_params import ProjectExportParams
 from labelbox.schema.media_type import MediaType
 from labelbox.schema.queue_mode import QueueMode
 from labelbox.schema.resource_tag import ResourceTag
+from labelbox.schema.task import Task
+from labelbox.schema.user import User
 
 if TYPE_CHECKING:
     from labelbox import BulkImportRequest
@@ -371,6 +375,71 @@ def _validate_datetime(string_date: str) -> bool:
                          self.uid)
             time.sleep(sleep_time)
 
+    """
+    Creates a project run export task with the given params and returns the task.
+    
+    >>>    export_task = export_v2("my_export_task", filter={"media_attributes": True})
+    
+    """
+
+    def export_v2(self, task_name: str,
+                  params: Optional[ProjectExportParams]) -> Task:
+        defaultParams: 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} }
+          """ % (mutation_name)
+        query_params = {
+            "input": {
+                "taskName": task_name,
+                "filters": {
+                    "projectId": self.uid
+                },
+                "params": {
+                    "includeAttachments":
+                        _params.get('attachments', False),
+                    "includeMediaAttributes":
+                        _params.get('media_attributes', False),
+                    "includeMetadata":
+                        _params.get('metadata_fields', False),
+                    "includeDataRowDetails":
+                        _params.get('data_row_details', False),
+                    "includeProjectDetails":
+                        _params.get('project_details', False),
+                    "includeLabels":
+                        _params.get('labels', False),
+                    "includePerformanceDetails":
+                        _params.get('performance_details', False),
+                },
+            }
+        }
+        res = self.client.execute(
+            create_task_query_str,
+            query_params,
+        )
+        res = res[mutation_name]
+        task_id = res["taskId"]
+        user: User = self.client.get_user()
+        tasks: List[Task] = list(
+            user.created_tasks(where=Entity.Task.uid == task_id))
+        # Cache user in a private variable as the relationship can't be
+        # resolved due to server-side limitations (see Task.created_by)
+        # for more info.
+        if len(tasks) != 1:
+            raise ResourceNotFoundError(Entity.Task, task_id)
+        task: Task = tasks[0]
+        task._user = user
+        return task
+
     def export_issues(self, status=None) -> str:
         """ Calls the server-side Issues exporting that
         returns the URL to that payload.

tests/integration/test_project.py

@@ -1,3 +1,4 @@
+import json
 import time
 import os
 
@@ -41,6 +42,49 @@ def test_project(client, rand_gen):
     assert project not in projects
 
 
+def test_project_export_v2(configured_project_with_label):
+    project, _, _, label = configured_project_with_label
+    label_id = label.uid
+    # Wait for exporter to retrieve latest labels
+    time.sleep(10)
+    task_name = "test_label_export_v2"
+
+    # TODO: Right now we don't have a way to test this
+    include_performance_details = True
+    task = project.export_v2(
+        task_name,
+        params={
+            "include_performance_details": include_performance_details,
+            "include_labels": True
+        })
+    assert task.name == task_name
+    task.wait_till_done()
+    assert task.status == "COMPLETE"
+
+    def download_result(result_url):
+        response = requests.get(result_url)
+        response.raise_for_status()
+        data = [json.loads(line) for line in response.text.splitlines()]
+        return data
+
+    task_results = download_result(task.result_url)
+
+    for task_result in task_results:
+        assert len(task_result['errors']) == 0
+        task_project = task_result['projects'][project.uid]
+        task_project_label_ids_set = set(
+            map(lambda prediction: prediction['id'], task_project['labels']))
+        assert label_id in task_project_label_ids_set
+
+        # TODO: Add back in when we have a way to test this
+        # if include_performance_details:
+        #     assert 'include_performance_details' in task_result and task_result[
+        #         'include_performance_details'] is not None
+        # else:
+        #     assert 'include_performance_details' not in task_result or task_result[
+        #         'include_performance_details'] is None
+
+
 def test_update_project_resource_tags(client, rand_gen):
 
     def delete_tag(tag_id: str):