Merge pull request #306 from Labelbox/gj/fix-number-type

Gareth · web-flow · commit f154a530a82e · 2021-10-10T10:31:33.000-07:00
Gj/fix number type
diff --git a/examples/basics/data_row_metadata.ipynb b/examples/basics/data_row_metadata.ipynb
@@ -374,8 +374,8 @@
     "                    value=dt,\n",
     "                ),\n",
     "                DataRowMetadataField(\n",
-    "                    schema_id=mdo.reserved_by_name[\"split\"].uid,\n",
-    "                    value=split\n",
+    "                    schema_id=split.parent,\n",
+    "                    value=split.uid\n",
     "                ),\n",
     "                DataRowMetadataField(\n",
     "                    schema_id=mdo.reserved_by_name[\"tag\"].uid,\n",
diff --git a/labelbox/schema/data_row_metadata.py b/labelbox/schema/data_row_metadata.py
@@ -1,6 +1,7 @@
 # type: ignore
 import datetime
 import warnings
+from copy import deepcopy
 from enum import Enum
 from itertools import chain
 from typing import List, Optional, Dict, Union, Callable, Type, Any, Generator
@@ -46,7 +47,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):
@@ -59,7 +62,7 @@ class Config:
 # Metadata base class
 class DataRowMetadataField(_CamelCaseMixin):
     schema_id: SchemaId
-    value: DataRowMetadataValue
+    value: Any
 
 
 class DataRowMetadata(_CamelCaseMixin):
@@ -85,7 +88,7 @@ class DataRowMetadataBatchResponse(_CamelCaseMixin):
 # Bulk upsert values
 class _UpsertDataRowMetadataInput(_CamelCaseMixin):
     schema_id: str
-    value: Union[str, List, dict]
+    value: Any
 
 
 # Batch of upsert values for a datarow
@@ -121,28 +124,30 @@ def __init__(self, client):
         self._batch_size = 50  # used for uploads and deletes
 
         self._raw_ontology = self._get_ontology()
+        self._build_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]):
@@ -151,7 +156,7 @@ def _make_name_index(fields: List[DataRowMetadataSchema]):
             if f.options:
                 index[f.name] = {}
                 for o in f.options:
-                    index[o.name] = o
+                    index[f.name][o.name] = o
             else:
                 index[f.name] = f
         return index
@@ -185,15 +190,17 @@ 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:
-            schema["uid"] = schema.pop("id")
+        copy = deepcopy(raw_ontology)
+        for schema in copy:
+            schema["uid"] = schema["id"]
             options = None
             if schema.get("options"):
                 options = []
                 for option in schema["options"]:
-                    option["uid"] = option.pop("id")
+                    option["uid"] = option["id"]
                     options.append(
                         DataRowMetadataSchema(**{
                             **option,
@@ -415,6 +422,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:
@@ -472,6 +481,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
diff --git a/tests/integration/test_data_row_metadata.py b/tests/integration/test_data_row_metadata.py
@@ -1,5 +1,5 @@
-from datetime import datetime
 import time
+from datetime import datetime
 
 import pytest
 
@@ -8,17 +8,29 @@
     DataRowMetadataOntology
 
 FAKE_SCHEMA_ID = "0" * 25
+FAKE_DATAROW_ID = "D" * 25
 SPLIT_SCHEMA_ID = "cko8sbczn0002h2dkdaxb5kal"
 TRAIN_SPLIT_ID = "cko8sbscr0003h2dk04w86hof"
 TEST_SPLIT_ID = "cko8scbz70005h2dkastwhgqt"
 EMBEDDING_SCHEMA_ID = "ckpyije740000yxdk81pbgjdc"
 TEXT_SCHEMA_ID = "cko8s9r5v0001h2dk9elqdidh"
 CAPTURE_DT_SCHEMA_ID = "cko8sdzv70006h2dk8jg64zvb"
 
+FAKE_NUMBER_FIELD = {
+    "id": FAKE_SCHEMA_ID,
+    "name": "number",
+    "kind": 'CustomMetadataNumber',
+    "reserved": False
+}
+
 
 @pytest.fixture
 def mdo(client):
-    yield client.get_data_row_metadata_ontology()
+    mdo = client.get_data_row_metadata_ontology()
+    mdo._raw_ontology = mdo._get_ontology()
+    mdo._raw_ontology.append(FAKE_NUMBER_FIELD)
+    mdo._build_ontology()
+    yield mdo
 
 
 @pytest.fixture
@@ -67,7 +79,21 @@ def make_metadata(dr_id) -> DataRowMetadata:
 def test_get_datarow_metadata_ontology(mdo):
     assert len(mdo.fields)
     assert len(mdo.reserved_fields)
-    assert len(mdo.custom_fields) == 0
+    assert len(mdo.custom_fields) == 1
+
+    split = mdo.reserved_by_name["split"]["train"]
+
+    assert DataRowMetadata(
+        data_row_id=FAKE_DATAROW_ID,
+        fields=[
+            DataRowMetadataField(
+                schema_id=mdo.reserved_by_name["captureDateTime"].uid,
+                value=datetime.utcnow(),
+            ),
+            DataRowMetadataField(schema_id=split.parent, value=split.uid),
+            DataRowMetadataField(schema_id=mdo.reserved_by_name["tag"].uid,
+                                 value="hello-world"),
+        ])
 
 
 def test_bulk_upsert_datarow_metadata(datarow, mdo: DataRowMetadataOntology):
@@ -127,7 +153,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)
@@ -217,23 +242,36 @@ def test_parse_raw_metadata(mdo):
     example = {
         'dataRowId':
             'ckr6kkfx801ui0yrtg9fje8xh',
-        'fields': [{
-            'schemaId': 'cko8s9r5v0001h2dk9elqdidh',
-            'value': 'my-new-message'
-        }, {
-            'schemaId': 'cko8sbczn0002h2dkdaxb5kal',
-            'value': {}
-        }, {
-            'schemaId': 'cko8sbscr0003h2dk04w86hof',
-            'value': {}
-        }, {
-            'schemaId': 'cko8sdzv70006h2dk8jg64zvb',
-            'value': '2021-07-20T21:41:14.606710Z'
-        }]
+        'fields': [
+            {
+                'schemaId': 'cko8s9r5v0001h2dk9elqdidh',
+                'value': 'my-new-message'
+            },
+            {
+                'schemaId': 'cko8sbczn0002h2dkdaxb5kal',
+                'value': {}
+            },
+            {
+                'schemaId': 'cko8sbscr0003h2dk04w86hof',
+                'value': {}
+            },
+            {
+                'schemaId': 'cko8sdzv70006h2dk8jg64zvb',
+                'value': '2021-07-20T21:41:14.606710Z'
+            },
+            {
+                'schemaId': FAKE_SCHEMA_ID,
+                '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)
