Merge pull request #791 from Labelbox/lgluszek/ontology-media-type

[CCV-1905][CCV-1906][CCV-1907] Pass media type when creating an ontology

Author: lgluszek

labelbox/client.py

@@ -36,7 +36,7 @@
 from labelbox.schema.slice import CatalogSlice
 from labelbox.schema.queue_mode import QueueMode
 
-from labelbox.schema.media_type import MediaType
+from labelbox.schema.media_type import MediaType, get_media_type_validation_error
 
 logger = logging.getLogger(__name__)
 
@@ -853,15 +853,18 @@ def rootSchemaPayloadToFeatureSchema(client, payload):
                                    rootSchemaPayloadToFeatureSchema,
                                    ['rootSchemaNodes', 'nextCursor'])
 
-    def create_ontology_from_feature_schemas(self, name,
-                                             feature_schema_ids) -> Ontology:
+    def create_ontology_from_feature_schemas(self,
+                                             name,
+                                             feature_schema_ids,
+                                             media_type=None) -> Ontology:
         """
         Creates an ontology from a list of feature schema ids
 
         Args:
             name (str): Name of the ontology
             feature_schema_ids (List[str]): List of feature schema ids corresponding to
                 top level tools and classifications to include in the ontology
+            media_type (MediaType or None): Media type of a new ontology
         Returns:
             The created Ontology
         """
@@ -891,9 +894,9 @@ def create_ontology_from_feature_schemas(self, name,
                     "Neither `tool` or `classification` found in the normalized feature schema"
                 )
         normalized = {'tools': tools, 'classifications': classifications}
-        return self.create_ontology(name, normalized)
+        return self.create_ontology(name, normalized, media_type)
 
-    def create_ontology(self, name, normalized) -> Ontology:
+    def create_ontology(self, name, normalized, media_type=None) -> Ontology:
         """
         Creates an ontology from normalized data
             >>> normalized = {"tools" : [{'tool': 'polygon',  'name': 'cat', 'color': 'black'}], "classifications" : []}
@@ -910,13 +913,27 @@ def create_ontology(self, name, normalized) -> Ontology:
         Args:
             name (str): Name of the ontology
             normalized (dict): A normalized ontology payload. See above for details.
+            media_type (MediaType or None): Media type of a new ontology
         Returns:
             The created Ontology
         """
+
+        if media_type:
+            if MediaType.is_supported(media_type):
+                media_type = media_type.value
+            else:
+                raise get_media_type_validation_error(media_type)
+
         query_str = """mutation upsertRootSchemaNodePyApi($data:  UpsertOntologyInput!){
                            upsertOntology(data: $data){ %s }
         } """ % query.results_query_part(Entity.Ontology)
-        params = {'data': {'name': name, 'normalized': json.dumps(normalized)}}
+        params = {
+            'data': {
+                'name': name,
+                'normalized': json.dumps(normalized),
+                'mediaType': media_type
+            }
+        }
         res = self.execute(query_str, params)
         return Entity.Ontology(self, res['upsertOntology'])
 
@@ -1035,9 +1052,9 @@ def _format_failed_rows(rows: Dict[str, str],
             )
 
         # Start assign global keys to data rows job
-        query_str = """mutation assignGlobalKeysToDataRowsPyApi($globalKeyDataRowLinks: [AssignGlobalKeyToDataRowInput!]!) { 
-            assignGlobalKeysToDataRows(data: {assignInputs: $globalKeyDataRowLinks}) { 
-                jobId 
+        query_str = """mutation assignGlobalKeysToDataRowsPyApi($globalKeyDataRowLinks: [AssignGlobalKeyToDataRowInput!]!) {
+            assignGlobalKeysToDataRows(data: {assignInputs: $globalKeyDataRowLinks}) {
+                jobId
             }
         }
         """
@@ -1172,7 +1189,7 @@ def _format_failed_rows(rows: List[str],
 
         # Query string for retrieving job status and result, if job is done
         result_query_str = """query getDataRowsForGlobalKeysResultPyApi($jobId: ID!) {
-            dataRowsForGlobalKeysResult(jobId: {id: $jobId}) { data { 
+            dataRowsForGlobalKeysResult(jobId: {id: $jobId}) { data {
                 fetchedDataRows { id }
                 notFoundGlobalKeys
                 accessDeniedGlobalKeys
@@ -1246,8 +1263,8 @@ def clear_global_keys(
 
             'Results' contains a list global keys that were successfully cleared.
 
-            'Errors' contains a list of global_keys correspond to the data rows that could not be 
-            modified, accessed by the user, or not found. 
+            'Errors' contains a list of global_keys correspond to the data rows that could not be
+            modified, accessed by the user, or not found.
         Examples:
             >>> job_result = client.get_data_row_ids_for_global_keys(["key1","key2"])
             >>> print(job_result['status'])
@@ -1271,7 +1288,7 @@ def _format_failed_rows(rows: List[str],
 
         # Query string for retrieving job status and result, if job is done
         result_query_str = """query clearGlobalKeysResultPyApi($jobId: ID!) {
-            clearGlobalKeysResult(jobId: {id: $jobId}) { data { 
+            clearGlobalKeysResult(jobId: {id: $jobId}) { data {
                 clearedGlobalKeys
                 failedToClearGlobalKeys
                 notFoundGlobalKeys

labelbox/schema/media_type.py

@@ -44,3 +44,9 @@ def get_supported_members(cls):
             item for item in cls.__members__
             if item not in ["Unknown", "Unsupported"]
         ]
+
+
+def get_media_type_validation_error(media_type):
+    return TypeError(f"{media_type} is not a valid media type. Use"
+                     f" any of {MediaType.get_supported_members()}"
+                     " from MediaType. Example: MediaType.Image.")