diff --git a/docarray/utils/create_dynamic_doc_class.py b/docarray/utils/create_dynamic_doc_class.py index d10f5bf23f..744fea58c3 100644 --- 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 index 9d9ec3d0b2..eba25911c4 100644 --- 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'