From 0d62ddf5e5b1620ef654dfad7fdbfc34771accc6 Mon Sep 17 00:00:00 2001
From: mnoszczak <mnoszczak@labelbox.com>
Date: Thu, 9 Mar 2023 17:31:57 +0100
Subject: [PATCH] Add media type override

---
 labelbox/schema/export_params.py  | 3 +++
 labelbox/schema/model_run.py      | 2 ++
 labelbox/schema/project.py        | 7 ++++++-
 tests/integration/test_project.py | 3 ++-
 4 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/labelbox/schema/export_params.py b/labelbox/schema/export_params.py
index fe72874d4..fa54ef81d 100644
--- a/labelbox/schema/export_params.py
+++ b/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):
diff --git a/labelbox/schema/model_run.py b/labelbox/schema/model_run.py
index 6bf993a37..55e70716f 100644
--- a/labelbox/schema/model_run.py
+++ b/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":
diff --git a/labelbox/schema/project.py b/labelbox/schema/project.py
index 99196fcf2..0858c9cf8 100644
--- a/labelbox/schema/project.py
+++ b/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":
diff --git a/tests/integration/test_project.py b/tests/integration/test_project.py
index bd3495af3..f985a6916 100644
--- a/tests/integration/test_project.py
+++ b/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(