fix: fix float in doc

JoanFM · JoanFM · commit 7730bbf04044 · 2024-03-18T11:21:35.000+01:00
diff --git a/docarray/utils/create_dynamic_doc_class.py b/docarray/utils/create_dynamic_doc_class.py
@@ -140,7 +140,9 @@ def _get_field_annotation_from_schema(
         for rec in range(num_recursions):
             ret = List[ret]
     elif field_type == 'number':
-        if num_recursions <= 1:
+        if num_recursions == 0:
+            ret = float
+        elif num_recursions == 1:
             # This is a hack because AnyTensor is more generic than a simple List and it comes as simple List
             if is_tensor:
                 ret = AnyTensor
diff --git a/tests/units/util/test_create_dynamic_code_class.py b/tests/units/util/test_create_dynamic_code_class.py
@@ -27,6 +27,8 @@ class Nested1Doc(BaseDoc):
     class CustomDoc(BaseDoc):
         tensor: Optional[AnyTensor] = None
         url: ImageUrl
+        num: float = 0.5
+        num_num: List[float] = [1.5, 2.5]
         lll: List[List[List[int]]] = [[[5]]]
         fff: List[List[List[float]]] = [[[5.2]]]
         single_text: TextDoc
@@ -47,6 +49,8 @@ class CustomDoc(BaseDoc):
     original_custom_docs = DocList[CustomDoc](
         [
             CustomDoc(
+                num=3.5,
+                num_num=[4.5, 5.5],
                 url='photo.jpg',
                 lll=[[[40]]],
                 fff=[[[40.2]]],
@@ -78,6 +82,8 @@ class CustomDoc(BaseDoc):
 
     assert len(custom_partial_da) == 1
     assert custom_partial_da[0].url == 'photo.jpg'
+    assert custom_partial_da[0].num == 3.5
+    assert custom_partial_da[0].num_num == [4.5, 5.5]
     assert custom_partial_da[0].lll == [[[40]]]
     if is_pydantic_v2:
         assert custom_partial_da[0].lu == [3, 4]
@@ -94,6 +100,8 @@ class CustomDoc(BaseDoc):
     assert custom_partial_da[0].single_text.text == 'single hey ha'
     assert custom_partial_da[0].single_text.embedding.shape == (2,)
     assert original_back[0].nested.nested.value == 'hello world'
+    assert original_back[0].num == 3.5
+    assert original_back[0].num_num == [4.5, 5.5]
     assert original_back[0].classvar == 'classvar'
     assert original_back[0].nested.classvar == 'classvar1'
     assert original_back[0].nested.nested.classvar == 'classvar2'
