FIX: parsing logic

Author: gdj0nes

labelbox/schema/data_row_metadata.py

@@ -50,7 +50,6 @@ def id(self):
 
 
 class _CamelCaseMixin(BaseModel):
-
     class Config:
         allow_population_by_field_name = True
         alias_generator = camel_case
@@ -158,7 +157,7 @@ def _make_name_index(fields: List[DataRowMetadataSchema]):
 
     @staticmethod
     def _make_id_index(
-        fields: List[DataRowMetadataSchema]
+            fields: List[DataRowMetadataSchema]
     ) -> Dict[SchemaId, DataRowMetadataSchema]:
         index = {}
         for f in fields:
@@ -198,7 +197,7 @@ def _parse_ontology(self) -> List[DataRowMetadataSchema]:
                         DataRowMetadataSchema(**{
                             **option,
                             **{
-                                "parent": schema["id"]
+                                "parent": schema["uid"]
                             }
                         }))
             schema["options"] = options
@@ -207,9 +206,9 @@ def _parse_ontology(self) -> List[DataRowMetadataSchema]:
         return fields
 
     def parse_metadata(
-        self, unparsed: List[Dict[str,
-                                  List[Union[str,
-                                             Dict]]]]) -> List[DataRowMetadata]:
+            self, unparsed: List[Dict[str,
+                                      List[Union[str,
+                                                 Dict]]]]) -> List[DataRowMetadata]:
         """ Parse metadata responses
 
         >>> mdo.parse_metadata([metdata])
@@ -270,7 +269,7 @@ def bulk_upsert(
             raise ValueError("Empty list passed")
 
         def _batch_upsert(
-            upserts: List[_UpsertBatchDataRowMetadata]
+                upserts: List[_UpsertBatchDataRowMetadata]
         ) -> List[DataRowMetadataBatchResponse]:
             query = """mutation UpsertDataRowMetadataBetaPyApi($metadata: [DataRowCustomMetadataBatchUpsertInput!]!) {
                 upsertDataRowCustomMetadata(data: $metadata){
@@ -303,13 +302,13 @@ def _batch_upsert(
                     fields=list(
                         chain.from_iterable(
                             self._parse_upsert(m) for m in m.fields))).dict(
-                                by_alias=True))
+                    by_alias=True))
 
         res = _batch_operations(_batch_upsert, items, self._batch_size)
         return res
 
     def bulk_delete(
-        self, deletes: List[DeleteDataRowMetadata]
+            self, deletes: List[DeleteDataRowMetadata]
     ) -> List[DataRowMetadataBatchResponse]:
         """ Delete metadata from a datarow by specifiying the fields you want to remove
 
@@ -336,7 +335,7 @@ def bulk_delete(
             raise ValueError("Empty list passed")
 
         def _batch_delete(
-            deletes: List[_DeleteBatchDataRowMetadata]
+                deletes: List[_DeleteBatchDataRowMetadata]
         ) -> List[DataRowMetadataBatchResponse]:
             query = """mutation DeleteDataRowMetadataBetaPyApi($deletes: [DataRowCustomMetadataBatchDeleteInput!]!) {
                 deleteDataRowCustomMetadata(data: $deletes) {
@@ -415,6 +414,8 @@ def _parse_upsert(
             parsed = _validate_parse_datetime(metadatum)
         elif schema.kind == DataRowMetadataKind.string:
             parsed = _validate_parse_text(metadatum)
+        elif schema.kind == DataRowMetadataKind.number:
+            parsed = _validate_parse_number(metadatum)
         elif schema.kind == DataRowMetadataKind.embedding:
             parsed = _validate_parse_embedding(metadatum)
         elif schema.kind == DataRowMetadataKind.enum:
@@ -455,9 +456,9 @@ def _batch_items(iterable: List[Any], size: int) -> Generator[Any, None, None]:
 
 
 def _batch_operations(
-    batch_function: _BatchFunction,
-    items: List,
-    batch_size: int = 100,
+        batch_function: _BatchFunction,
+        items: List,
+        batch_size: int = 100,
 ):
     response = []
 
@@ -472,6 +473,12 @@ def _validate_parse_embedding(
     return [field.dict(by_alias=True)]
 
 
+def _validate_parse_number(
+        field: DataRowMetadataField
+) -> List[Dict[str, Union[SchemaId, Number]]]:
+    return [field.dict(by_alias=True)]
+
+
 def _validate_parse_datetime(
         field: DataRowMetadataField) -> List[Dict[str, Union[SchemaId, str]]]:
     # TODO: better validate tzinfo

tests/integration/test_data_row_metadata.py

@@ -1,5 +1,5 @@
-from datetime import datetime
 import time
+from datetime import datetime
 
 import pytest
 
@@ -24,11 +24,11 @@ def mdo(client):
 @pytest.fixture
 def big_dataset(dataset: Dataset, image_url):
     task = dataset.create_data_rows([
-        {
-            "row_data": image_url,
-            "external_id": "my-image"
-        },
-    ] * 250)
+                                        {
+                                            "row_data": image_url,
+                                            "external_id": "my-image"
+                                        },
+                                    ] * 250)
     task.wait_till_done()
 
     yield dataset
@@ -38,11 +38,11 @@ def big_dataset(dataset: Dataset, image_url):
 def wait_for_embeddings_svc(data_row_ids, mdo):
     for idx in range(5):
         if all([
-                len(metadata.fields)
-                for metadata in mdo.bulk_export(data_row_ids)
+            len(metadata.fields)
+            for metadata in mdo.bulk_export(data_row_ids)
         ]):
             return
-        time.sleep((idx + 1)**2)
+        time.sleep((idx + 1) ** 2)
     raise Exception("Embedding svc failed to update metadata.")
 
 
@@ -127,7 +127,6 @@ def test_bulk_partial_delete_datarow_metadata(datarow, mdo):
 
 
 def test_large_bulk_delete_datarow_metadata(big_dataset, mdo):
-
     metadata = []
     data_row_ids = [dr.uid for dr in big_dataset.data_rows()]
     wait_for_embeddings_svc(data_row_ids, mdo)
@@ -229,11 +228,19 @@ def test_parse_raw_metadata(mdo):
         }, {
             'schemaId': 'cko8sdzv70006h2dk8jg64zvb',
             'value': '2021-07-20T21:41:14.606710Z'
-        }]
+        }, {
+            'schemaId': 'cko8sdzv70006h2dk8jg64zvb',
+            'value': 0.5
+        },
+        ]
     }
 
     parsed = mdo.parse_metadata([example])
     assert len(parsed) == 1
-    row = parsed[0]
-    assert row.data_row_id == example["dataRowId"]
-    assert len(row.fields) == 3
+    for row in parsed:
+        assert row.data_row_id == example["dataRowId"]
+        assert len(row.fields) == 4
+
+    for row in parsed:
+        for field in row.fields:
+            assert mdo._parse_upsert(field)