Add MediaType support to Project

Author: paulnoirel

labelbox/client.py

@@ -611,6 +611,20 @@ def create_project(self, **kwargs) -> Project:
             InvalidAttributeError: If the Project type does not contain
                 any of the attribute names given in kwargs.
         """
+        media_type = kwargs.get("media_type")
+        if media_type:
+            if isinstance(media_type, Project.MediaType
+                         ) and media_type is not Project.MediaType.Unknown:
+                kwargs["media_type"] = media_type.value
+            else:
+                media_types = [
+                    item for item in Project.MediaType.__members__
+                    if item != "Unknown"
+                ]
+                raise TypeError(
+                    f"{media_type} is not a supported type. Please use any of {media_types} from the {type(media_type).__name__} enumeration."
+                )
+
         return self._create(Entity.Project, kwargs)
 
     def get_roles(self) -> List[Role]:

labelbox/orm/db_object.py

@@ -69,7 +69,7 @@ def _set_field_values(self, field_values):
                         "Failed to convert value '%s' to datetime for "
                         "field %s", value, field)
             elif isinstance(field.field_type, Field.EnumType):
-                value = field.field_type.enum_cls[value]
+                value = field.field_type.enum_cls(value)
             setattr(self, field.name, value)
 
     def __repr__(self):

labelbox/schema/project.py

@@ -1,4 +1,4 @@
-import enum
+from enum import Enum
 import json
 import logging
 import time
@@ -63,6 +63,37 @@ class Project(DbObject, Updateable, Deletable):
         benchmarks (Relationship): `ToMany` relationship to Benchmark
         ontology (Relationship): `ToOne` relationship to Ontology
     """
+
+    class MediaType(Enum):
+        """add DOCUMENT, GEOSPATIAL_TILE, SIMPLE_TILE to match the UI choices"""
+        Audio = "AUDIO"
+        Conversational = "CONVERSATIONAL"
+        Dicom = "DICOM"
+        Document = "PDF"
+        Geospatial_Tile = "TMS_GEO"
+        Image = "IMAGE"
+        Json = "JSON"
+        Pdf = "PDF"
+        Simple_Tile = "TMS_SIMPLE"
+        Text = "TEXT"
+        Tms_Geo = "TMS_GEO"
+        Tms_Simple = "TMS_SIMPLE"
+        Video = "VIDEO"
+        Unknown = "UNKNOWN"
+
+        @classmethod
+        def _missing_(cls, name):
+            """Handle missing null data types for projects 
+            created without setting allowedMediaType"""
+            # return Project.MediaType.UNKNOWN
+
+            if name is None:
+                return cls.Unknown
+
+            for member in cls.__members__:
+                if member.name == name.upper():
+                    return member
+
     name = Field.String("name")
     description = Field.String("description")
     updated_at = Field.DateTime("updated_at")
@@ -71,6 +102,8 @@ class Project(DbObject, Updateable, Deletable):
     last_activity_time = Field.DateTime("last_activity_time")
     auto_audit_number_of_labels = Field.Int("auto_audit_number_of_labels")
     auto_audit_percentage = Field.Float("auto_audit_percentage")
+    # Bind data_type and allowedMediaTYpe using the GraphQL type MediaType
+    media_type = Field.Enum(MediaType, "media_type", "allowedMediaType")
 
     # Relationships
     datasets = Relationship.ToMany("Dataset", True)
@@ -85,7 +118,7 @@ class Project(DbObject, Updateable, Deletable):
     benchmarks = Relationship.ToMany("Benchmark", False)
     ontology = Relationship.ToOne("Ontology", True)
 
-    class QueueMode(enum.Enum):
+    class QueueMode(Enum):
         Batch = "Batch"
         Dataset = "Dataset"
 
@@ -94,6 +127,20 @@ def update(self, **kwargs):
         if mode:
             self._update_queue_mode(mode)
 
+        media_type = kwargs.get("media_type")
+        if media_type:
+            if isinstance(media_type, Project.MediaType
+                         ) and media_type != Project.MediaType.Unknown:
+                kwargs["media_type"] = media_type.value
+            else:
+                media_types = [
+                    item for item in Project.MediaType.__members__
+                    if item != "Unknown"
+                ]
+                raise TypeError(
+                    f"{media_type} is not a supported type. Please use any of {media_types} from the {type(media_type).__name__} enumeration."
+                )
+
         return super().update(**kwargs)
 
     def members(self) -> PaginatedCollection:

tests/integration/test_project.py

@@ -1,4 +1,3 @@
-import json
 import time
 import os
 
@@ -44,11 +43,19 @@ def test_project(client, rand_gen):
     assert set(final) == set(before)
 
 
-@pytest.mark.skip(
-    reason="this will fail if run multiple times, limit is defaulted to 3 per org"
-    "add this back in when either all test orgs have unlimited, or we delete all tags befoer running"
-)
 def test_update_project_resource_tags(client, rand_gen):
+
+    def delete_tag(tag_id: str):
+        """Deletes a tag given the tag uid. Currently internal use only so this is not public"""
+        res = client.execute(
+            """mutation deleteResourceTagPyApi($tag_id: String!) {
+        deleteResourceTag(input: {id: $tag_id}) {
+            id
+        }
+        }
+        """, {"tag_id": tag_id})
+        return res
+
     before = list(client.get_projects())
     for o in before:
         assert isinstance(o, Project)
@@ -92,6 +99,9 @@ def test_update_project_resource_tags(client, rand_gen):
     assert len(project_resource_tag) == 1
     assert project_resource_tag[0].uid == tagA.uid
 
+    delete_tag(tagA.uid)
+    delete_tag(tagB.uid)
+
 
 def test_project_filtering(client, rand_gen):
     name_1 = rand_gen(str)
@@ -191,3 +201,25 @@ def test_queue_mode(configured_project: Project):
     ) == configured_project.QueueMode.Dataset
     configured_project.update(queue_mode=configured_project.QueueMode.Batch)
     assert configured_project.queue_mode() == configured_project.QueueMode.Batch
+
+
+def test_media_type(client, configured_project: Project, rand_gen):
+    # Existing project
+    assert configured_project.media_type is None or isinstance(
+        configured_project.media_type, Project.MediaType)
+
+    # No media_type
+    project = client.create_project(name=rand_gen(str))
+    assert project.media_type == Project.MediaType.Unknown
+    project.update(media_type=Project.MediaType.Image)
+    assert project.media_type == Project.MediaType.Image
+    project.delete()
+
+    for media_type in Project.MediaType:
+        if media_type == Project.MediaType.Unknown:
+            continue
+
+        project = client.create_project(name=rand_gen(str),
+                                        media_type=media_type)
+        assert project.media_type == media_type
+        project.delete()