Add catalog slice SDK support

Author: mnoszczak

labelbox/schema/export_params.py

@@ -1,6 +1,6 @@
 import sys
 
-from typing import Optional
+from typing import Optional, List
 
 from labelbox.schema.media_type import MediaType
 if sys.version_info >= (3, 8):
@@ -22,6 +22,15 @@ class ProjectExportParams(DataRowParams):
     performance_details: Optional[bool]
 
 
+class CatalogSliceExportParams(DataRowParams):
+    project_details: Optional[bool]
+    label_details: Optional[bool]
+    performance_details: Optional[bool]
+    model_runs_ids: Optional[List[str]]
+    projects_ids: Optional[List[str]]
+    pass
+
+
 class ModelRunExportParams(DataRowParams):
     # TODO: Add model run fields
     pass

labelbox/schema/slice.py

@@ -1,6 +1,11 @@
+from typing import Optional, List
+from labelbox.exceptions import ResourceNotFoundError
 from labelbox.orm.db_object import DbObject
-from labelbox.orm.model import Field
+from labelbox.orm.model import Entity, Field
 from labelbox.pagination import PaginatedCollection
+from labelbox.schema.export_params import CatalogSliceExportParams
+from labelbox.schema.task import Task
+from labelbox.schema.user import User
 
 
 class Slice(DbObject):
@@ -59,6 +64,85 @@ def get_data_row_ids(self) -> PaginatedCollection:
             obj_class=lambda _, data_row_id: data_row_id,
             cursor_path=['getDataRowIdsBySavedQuery', 'pageInfo', 'endCursor'])
 
+    def export_v2(self,
+                  task_name: Optional[str] = None,
+                  params: Optional[CatalogSliceExportParams] = None) -> Task:
+        """
+        Creates a slice export task with the given params and returns the task.
+        >>>     slice = client.get_catalog_slice("SLICE_ID")
+        >>>     task = slice.export_v2(
+        >>>         params={"performance_details": False, "label_details": True}
+        >>>     )
+        >>>     task.wait_till_done()
+        >>>     task.result
+        """
+
+        _params = params or CatalogSliceExportParams({
+            "attachments": False,
+            "metadata_fields": False,
+            "data_row_details": False,
+            "project_details": False,
+            "performance_details": False,
+            "label_details": False,
+            "media_type_override": None,
+            "model_runs_ids": None,
+            "projects_ids": None,
+        })
+
+        mutation_name = "exportDataRowsInSlice"
+        create_task_query_str = """mutation exportDataRowsInSlicePyApi($input: ExportDataRowsInSliceInput!){
+          %s(input: $input) {taskId} }
+          """ % (mutation_name)
+
+        media_type_override = _params.get('media_type_override', None)
+        query_params = {
+            "input": {
+                "taskName": task_name,
+                "filters": {
+                    "sliceId": self.uid
+                },
+                "params": {
+                    "mediaTypeOverride":
+                        media_type_override.value
+                        if media_type_override is not None else None,
+                    "includeAttachments":
+                        _params.get('attachments', False),
+                    "includeMetadata":
+                        _params.get('metadata_fields', False),
+                    "includeDataRowDetails":
+                        _params.get('data_row_details', False),
+                    "includeProjectDetails":
+                        _params.get('project_details', False),
+                    "includePerformanceDetails":
+                        _params.get('performance_details', False),
+                    "includeLabelDetails":
+                        _params.get('label_details', False),
+                    "projectIds":
+                        _params.get('projects_ids', None),
+                    "modelRunIds":
+                        _params.get('model_runs_ids', None),
+                },
+            }
+        }
+
+        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
+
 
 class ModelSlice(Slice):
     """