Vb/remove label data plt 37 (#1527)

Author: vbrodsky

labelbox/data/annotation_types/data/generic_data_row_data.py

@@ -0,0 +1,26 @@
+from typing import Callable, Literal, Optional
+
+from labelbox import pydantic_compat
+from labelbox.data.annotation_types.data.base_data import BaseData
+from labelbox.utils import _NoCoercionMixin
+
+
+class GenericDataRowData(BaseData, _NoCoercionMixin):
+    """Generic data row data. This is replacing all other DataType passed into Label
+    """
+    url: Optional[str] = None
+    class_name: Literal["GenericDataRowData"] = "GenericDataRowData"
+
+    def create_url(self, signer: Callable[[bytes], str]) -> Optional[str]:
+        return self.url
+
+    @pydantic_compat.root_validator(pre=True)
+    def validate_one_datarow_key_present(cls, data):
+        keys = ['external_id', 'global_key', 'uid']
+        count = sum([key in data for key in keys])
+
+        if count < 1:
+            raise ValueError(f"Exactly one of {keys} must be present.")
+        if count > 1:
+            raise ValueError(f"Only one of {keys} can be present.")
+        return data

labelbox/data/annotation_types/label.py

@@ -5,12 +5,13 @@
 from labelbox import pydantic_compat
 
 import labelbox
+from labelbox.data.annotation_types.data.generic_data_row_data import GenericDataRowData
 from labelbox.data.annotation_types.data.tiled_image import TiledImageData
 from labelbox.schema import ontology
 from .annotation import ClassificationAnnotation, ObjectAnnotation
 from .relationship import RelationshipAnnotation
 from .classification import ClassificationAnswer
-from .data import AudioData, ConversationData, DicomData, DocumentData, HTMLData, ImageData, MaskData, TextData, VideoData, LlmPromptCreationData, LlmPromptResponseCreationData, LlmResponseCreationData
+from .data import AudioData, ConversationData, DicomData, DocumentData, HTMLData, ImageData, TextData, VideoData, LlmPromptCreationData, LlmPromptResponseCreationData, LlmResponseCreationData
 from .geometry import Mask
 from .metrics import ScalarMetric, ConfusionMatrixMetric
 from .types import Cuid
@@ -21,14 +22,14 @@
 DataType = Union[VideoData, ImageData, TextData, TiledImageData, AudioData,
                  ConversationData, DicomData, DocumentData, HTMLData,
                  LlmPromptCreationData, LlmPromptResponseCreationData,
-                 LlmResponseCreationData]
+                 LlmResponseCreationData, GenericDataRowData]
 
 
 class Label(pydantic_compat.BaseModel):
     """Container for holding data and annotations
 
     >>> Label(
-    >>>    data = ImageData(url = "http://my-img.jpg"),
+    >>>    data = {'global_key': 'my-data-row-key'} # also accepts uid, external_id as keys
     >>>    annotations = [
     >>>        ObjectAnnotation(
     >>>            value = Point(x = 10, y = 10),
@@ -39,7 +40,8 @@ class Label(pydantic_compat.BaseModel):
 
     Args:
         uid: Optional Label Id in Labelbox
-        data: Data of Label, Image, Video, Text
+        data: Data of Label, Image, Video, Text or dict with a single key uid | global_key | external_id. 
+            Note use of classes as data is deprecated. Use GenericDataRowData or dict with a single key instead.
         annotations: List of Annotations in the label
         extra: additional context
     """
@@ -51,6 +53,16 @@ class Label(pydantic_compat.BaseModel):
                             RelationshipAnnotation]] = []
     extra: Dict[str, Any] = {}
 
+    @pydantic_compat.root_validator(pre=True)
+    def validate_data(cls, label):
+        if isinstance(label.get("data"), Dict):
+            label["data"]["class_name"] = "GenericDataRowData"
+        else:
+            warnings.warn(
+                f"Using {type(label['data']).__name__} class for label.data is deprecated. "
+                "Use a dict or an instance of GenericDataRowData instead.")
+        return label
+
     def object_annotations(self) -> List[ObjectAnnotation]:
         return self._get_annotations_by_type(ObjectAnnotation)
 

labelbox/schema/id_type.py

@@ -1,7 +1,18 @@
-from strenum import StrEnum
+import sys
 
+if sys.version_info >= (3, 9):
+    from strenum import StrEnum
 
-class IdType(StrEnum):
+    class BaseStrEnum(StrEnum):
+        pass
+else:
+    from enum import Enum
+
+    class BaseStrEnum(str, Enum):
+        pass
+
+
+class IdType(BaseStrEnum):
     """
     The type of id used to identify a data row.
     

tests/conftest.py

@@ -121,13 +121,13 @@ def rest_url(environ: str) -> str:
     return 'http://host.docker.internal:8080/api/v1'
 
 
-def testing_api_key(environ: str) -> str:
-    for var in [
-            "LABELBOX_TEST_API_KEY_PROD", "LABELBOX_TEST_API_KEY_STAGING",
-            "LABELBOX_TEST_API_KEY_CUSTOM", "LABELBOX_TEST_API_KEY_LOCAL",
-            "LABELBOX_TEST_API_KEY"
-    ]:
-        value = os.environ.get(var)
+def testing_api_key(environ: Environ) -> str:
+    keys = [
+        f"LABELBOX_TEST_API_KEY_{environ.value.upper()}",
+        "LABELBOX_TEST_API_KEY"
+    ]
+    for key in keys:
+        value = os.environ.get(key)
         if value is not None:
             return value
     raise Exception("Cannot find API to use for tests")
@@ -147,7 +147,6 @@ def __init__(self, environ: str) -> None:
         api_url = graphql_url(environ)
         api_key = testing_api_key(environ)
         rest_endpoint = rest_url(environ)
-
         super().__init__(api_key,
                          api_url,
                          enable_experimental=True,

tests/data/annotation_import/conftest.py

@@ -1885,3 +1885,64 @@ def bbox_video_annotation_objects():
     ]
 
     return bbox_annotation
+
+
+class Helpers:
+
+    @staticmethod
+    def remove_keys_recursive(d, keys):
+        for k in keys:
+            if k in d:
+                del d[k]
+        for k, v in d.items():
+            if isinstance(v, dict):
+                Helpers.remove_keys_recursive(v, keys)
+            elif isinstance(v, list):
+                for i in v:
+                    if isinstance(i, dict):
+                        Helpers.remove_keys_recursive(i, keys)
+
+    @staticmethod
+    # NOTE this uses quite a primitive check for cuids but I do not think it is worth coming up with a better one
+    # Also this function is NOT written with performance in mind, good for small to mid size dicts like we have in our test
+    def rename_cuid_key_recursive(d):
+        new_key = "<cuid>"
+        for k in list(d.keys()):
+            if len(k) == 25 and not k.isalpha():  # primitive check for cuid
+                d[new_key] = d.pop(k)
+        for k, v in d.items():
+            if isinstance(v, dict):
+                Helpers.rename_cuid_key_recursive(v)
+            elif isinstance(v, list):
+                for i in v:
+                    if isinstance(i, dict):
+                        Helpers.rename_cuid_key_recursive(i)
+
+    @staticmethod
+    def set_project_media_type_from_data_type(project, data_type_class):
+
+        def to_pascal_case(name: str) -> str:
+            return "".join([word.capitalize() for word in name.split("_")])
+
+        data_type_string = data_type_class.__name__[:-4].lower()
+        media_type = to_pascal_case(data_type_string)
+        if media_type == "Conversation":
+            media_type = "Conversational"
+        elif media_type == "Llmpromptcreation":
+            media_type = "LLMPromptCreation"
+        elif media_type == "Llmpromptresponsecreation":
+            media_type = "LLMPromptResponseCreation"
+        elif media_type == "Llmresponsecreation":
+            media_type = "Text"
+        elif media_type == "Genericdatarow":
+            media_type = "Image"
+        project.update(media_type=MediaType[media_type])
+
+    @staticmethod
+    def find_data_row_filter(data_row):
+        return lambda dr: dr['data_row']['id'] == data_row.uid
+
+
+@pytest.fixture
+def helpers():
+    return Helpers

tests/data/annotation_import/test_data_types.py

@@ -5,7 +5,6 @@
 
 import labelbox as lb
 from labelbox.data.annotation_types.data.video import VideoData
-from labelbox.schema.data_row import DataRow
 from labelbox.schema.media_type import MediaType
 import labelbox.types as lb_types
 from labelbox.data.annotation_types.data import (
@@ -70,35 +69,6 @@
 ]
 
 
-def remove_keys_recursive(d, keys):
-    for k in keys:
-        if k in d:
-            del d[k]
-    for k, v in d.items():
-        if isinstance(v, dict):
-            remove_keys_recursive(v, keys)
-        elif isinstance(v, list):
-            for i in v:
-                if isinstance(i, dict):
-                    remove_keys_recursive(i, keys)
-
-
-# NOTE this uses quite a primitive check for cuids but I do not think it is worth coming up with a better one
-# Also this function is NOT written with performance in mind, good for small to mid size dicts like we have in our test
-def rename_cuid_key_recursive(d):
-    new_key = "<cuid>"
-    for k in list(d.keys()):
-        if len(k) == 25 and not k.isalpha():  # primitive check for cuid
-            d[new_key] = d.pop(k)
-    for k, v in d.items():
-        if isinstance(v, dict):
-            rename_cuid_key_recursive(v)
-        elif isinstance(v, list):
-            for i in v:
-                if isinstance(i, dict):
-                    rename_cuid_key_recursive(i)
-
-
 def get_annotation_comparison_dicts_from_labels(labels):
     labels_ndjson = list(NDJsonConverter.serialize(labels))
     for annotation in labels_ndjson:
@@ -198,12 +168,13 @@ def test_import_data_types(
     data_row_json_by_data_type,
     annotations_by_data_type,
     data_type_class,
+    helpers,
 ):
     project = configured_project
     project_id = project.uid
     dataset = initial_dataset
 
-    set_project_media_type_from_data_type(project, data_type_class)
+    helpers.set_project_media_type_from_data_type(project, data_type_class)
 
     data_type_string = data_type_class.__name__[:-4].lower()
     data_row_ndjson = data_row_json_by_data_type[data_type_string]
@@ -241,12 +212,13 @@ def test_import_data_types_by_global_key(
     rand_gen,
     data_row_json_by_data_type,
     annotations_by_data_type,
+    helpers,
 ):
     project = configured_project
     project_id = project.uid
     dataset = initial_dataset
     data_type_class = ImageData
-    set_project_media_type_from_data_type(project, data_type_class)
+    helpers.set_project_media_type_from_data_type(project, data_type_class)
 
     data_row_ndjson = data_row_json_by_data_type["image"]
     data_row_ndjson["global_key"] = str(uuid.uuid4())
@@ -287,24 +259,6 @@ def validate_iso_format(date_string: str):
     assert parsed_t.second is not None
 
 
-def to_pascal_case(name: str) -> str:
-    return "".join([word.capitalize() for word in name.split("_")])
-
-
-def set_project_media_type_from_data_type(project, data_type_class):
-    data_type_string = data_type_class.__name__[:-4].lower()
-    media_type = to_pascal_case(data_type_string)
-    if media_type == "Conversation":
-        media_type = "Conversational"
-    elif media_type == "Llmpromptcreation":
-        media_type = "LLMPromptCreation"
-    elif media_type == "Llmpromptresponsecreation":
-        media_type = "LLMPromptResponseCreation"
-    elif media_type == "Llmresponsecreation":
-        media_type = "Text"
-    project.update(media_type=MediaType[media_type])
-
-
 @pytest.mark.parametrize(
     "data_type_class",
     [
@@ -331,12 +285,13 @@ def test_import_data_types_v2(
     exports_v2_by_data_type,
     export_v2_test_helpers,
     rand_gen,
+    helpers,
 ):
     project = configured_project
     dataset = initial_dataset
     project_id = project.uid
 
-    set_project_media_type_from_data_type(project, data_type_class)
+    helpers.set_project_media_type_from_data_type(project, data_type_class)
 
     data_type_string = data_type_class.__name__[:-4].lower()
     data_row_ndjson = data_row_json_by_data_type[data_type_string]
@@ -381,9 +336,9 @@ def test_import_data_types_v2(
     exported_project_labels = exported_project["labels"][0]
     exported_annotations = exported_project_labels["annotations"]
 
-    remove_keys_recursive(exported_annotations,
-                          ["feature_id", "feature_schema_id"])
-    rename_cuid_key_recursive(exported_annotations)
+    helpers.remove_keys_recursive(exported_annotations,
+                                  ["feature_id", "feature_schema_id"])
+    helpers.rename_cuid_key_recursive(exported_annotations)
     assert exported_annotations == exports_v2_by_data_type[data_type_string]
 
     data_row = client.get_data_row(data_row.uid)
@@ -400,10 +355,11 @@ def test_import_label_annotations(
     data_class,
     annotations,
     rand_gen,
+    helpers,
 ):
     project = configured_project_with_one_data_row
     dataset = initial_dataset
-    set_project_media_type_from_data_type(project, data_class)
+    helpers.set_project_media_type_from_data_type(project, data_class)
 
     data_row_json = data_row_json_by_data_type[data_type]
     data_row = create_data_row_for_project(project, dataset, data_row_json,
@@ -471,10 +427,11 @@ def test_import_mal_annotations(
     annotations,
     rand_gen,
     one_datarow,
+    helpers,
 ):
     data_row = one_datarow
-    set_project_media_type_from_data_type(configured_project_with_one_data_row,
-                                          data_class)
+    helpers.set_project_media_type_from_data_type(
+        configured_project_with_one_data_row, data_class)
 
     configured_project_with_one_data_row.create_batch(
         rand_gen(str),
@@ -500,12 +457,13 @@ def test_import_mal_annotations(
 
 def test_import_mal_annotations_global_key(client,
                                            configured_project_with_one_data_row,
-                                           rand_gen, one_datarow_global_key):
+                                           rand_gen, one_datarow_global_key,
+                                           helpers):
     data_class = lb_types.VideoData
     data_row = one_datarow_global_key
     annotations = [video_mask_annotation]
-    set_project_media_type_from_data_type(configured_project_with_one_data_row,
-                                          data_class)
+    helpers.set_project_media_type_from_data_type(
+        configured_project_with_one_data_row, data_class)
 
     configured_project_with_one_data_row.create_batch(
         rand_gen(str),