FIX: type conversion

Author: gdj0nes

labelbox/schema/data_row_metadata.py

@@ -46,7 +46,9 @@ def id(self):
 OptionId: Type[SchemaId] = SchemaId  # enum option
 Number: Type[float] = float
 
-DataRowMetadataValue = Union[Embedding, DateTime, String, OptionId, Number]
+DataRowMetadataValue = Union[Embedding, Number, DateTime, String, OptionId]
+# primitives used in uploads
+_DataRowMetadataValuePrimitives = Union[str, List, dict, float]
 
 
 class _CamelCaseMixin(BaseModel):
@@ -84,7 +86,7 @@ class DataRowMetadataBatchResponse(_CamelCaseMixin):
 # Bulk upsert values
 class _UpsertDataRowMetadataInput(_CamelCaseMixin):
     schema_id: str
-    value: Union[str, List, dict]
+    value: _DataRowMetadataValuePrimitives
 
 
 # Batch of upsert values for a datarow
@@ -121,27 +123,24 @@ def __init__(self, client):
 
         self._raw_ontology = self._get_ontology()
 
+    def _build_ontology(self):
         # all fields
-        self.fields = self._parse_ontology()
+        self.fields = self._parse_ontology(self._raw_ontology)
         self.fields_by_id = self._make_id_index(self.fields)
 
         # reserved fields
         self.reserved_fields: List[DataRowMetadataSchema] = [
             f for f in self.fields if f.reserved
         ]
         self.reserved_by_id = self._make_id_index(self.reserved_fields)
-        self.reserved_by_name: Dict[str, DataRowMetadataSchema] = {
-            f.name: f for f in self.reserved_fields
-        }
+        self.reserved_by_name: Dict[str, DataRowMetadataSchema] = self._make_name_index(self.reserved_fields)
 
         # custom fields
         self.custom_fields: List[DataRowMetadataSchema] = [
             f for f in self.fields if not f.reserved
         ]
         self.custom_by_id = self._make_id_index(self.custom_fields)
-        self.custom_by_name: Dict[str, DataRowMetadataSchema] = {
-            f.name: f for f in self.custom_fields
-        }
+        self.custom_by_name: Dict[str, DataRowMetadataSchema] = self._make_name_index(self.custom_fields)
 
     @staticmethod
     def _make_name_index(fields: List[DataRowMetadataSchema]):
@@ -184,9 +183,10 @@ def _get_ontology(self) -> List[Dict[str, Any]]:
         """
         return self._client.execute(query)["customMetadataOntology"]
 
-    def _parse_ontology(self) -> List[DataRowMetadataSchema]:
+    @staticmethod
+    def _parse_ontology(raw_ontology) -> List[DataRowMetadataSchema]:
         fields = []
-        for schema in self._raw_ontology:
+        for schema in raw_ontology:
             schema["uid"] = schema.pop("id")
             options = None
             if schema.get("options"):

tests/integration/test_data_row_metadata.py

@@ -15,10 +15,34 @@
 TEXT_SCHEMA_ID = "cko8s9r5v0001h2dk9elqdidh"
 CAPTURE_DT_SCHEMA_ID = "cko8sdzv70006h2dk8jg64zvb"
 
+FAKE_NUMBER_FIELD = {
+    "id": FAKE_SCHEMA_ID,
+    "name": "number",
+    "kind": 'CustomMetadataNumber',
+    "reserved": False
+}
+
+"""
+customMetadataOntology {
+                id
+                name
+                kind
+                reserved
+                options {
+                    id
+                    kind
+                    name
+                    reserved
+                }
+        }}"""
+
 
 @pytest.fixture
 def mdo(client):
-    yield client.get_data_row_metadata_ontology()
+    mdo = client.get_data_row_metadata_ontology()
+    mdo._raw_ontology.append(FAKE_NUMBER_FIELD)
+    mdo._build_ontology()
+    yield mdo
 
 
 @pytest.fixture
@@ -229,7 +253,7 @@ def test_parse_raw_metadata(mdo):
             'schemaId': 'cko8sdzv70006h2dk8jg64zvb',
             'value': '2021-07-20T21:41:14.606710Z'
         }, {
-            'schemaId': 'cko8sdzv70006h2dk8jg64zvb',
+            'schemaId': FAKE_SCHEMA_ID,
             'value': 0.5
         },
         ]
@@ -243,4 +267,4 @@ def test_parse_raw_metadata(mdo):
 
     for row in parsed:
         for field in row.fields:
-            assert mdo._parse_upsert(field)
+            assert mdo._parse_upsert(field)