Add media type override

Author: mnoszczak

labelbox/schema/export_params.py

@@ -1,6 +1,8 @@
 import sys
 
 from typing import Optional
+
+from labelbox.schema.media_type import MediaType
 if sys.version_info >= (3, 8):
     from typing import TypedDict
 else:
@@ -11,6 +13,7 @@ class DataRowParams(TypedDict):
     data_row_details: Optional[bool]
     metadata_fields: Optional[bool]
     attachments: Optional[bool]
+    media_type_override: Optional[MediaType]
 
 
 class ProjectExportParams(DataRowParams):

labelbox/schema/model_run.py

@@ -521,6 +521,8 @@ def export_v2(self,
                     "modelRunId": self.uid
                 },
                 "params": {
+                    "mediaTypeOverride":
+                        _params.get('media_type_override', None),
                     "includeAttachments":
                         _params.get('attachments', False),
                     "includeMetadata":

labelbox/schema/project.py

@@ -430,7 +430,8 @@ def export_v2(self,
             "data_row_details": False,
             "project_details": False,
             "performance_details": False,
-            "label_details": False
+            "label_details": False,
+            "media_type_override": None
         })
 
         _filters = filters or ProjectExportFilters({
@@ -451,6 +452,7 @@ def _get_timezone() -> str:
           """ % (mutation_name)
 
         search_query: List[Dict[str, Collection[str]]] = []
+        media_type_override = _params.get('media_type_override', None)
         query_params = {
             "input": {
                 "taskName": task_name,
@@ -462,6 +464,9 @@ def _get_timezone() -> str:
                     }
                 },
                 "params": {
+                    "mediaTypeOverride":
+                        media_type_override.value
+                        if media_type_override is not None else None,
                     "includeAttachments":
                         _params.get('attachments', False),
                     "includeMetadata":

tests/integration/test_project.py

@@ -59,7 +59,8 @@ def test_project_export_v2(configured_project_with_label):
         },
         params={
             "include_performance_details": include_performance_details,
-            "include_labels": True
+            "include_labels": True,
+            "media_type_override": MediaType.Image
         })
 
     task_to = project.export_v2(