test: fix more tests

Signed-off-by: Joan Fontanals Martinez <joan.martinez@jina.ai>
docarray · JoanFM · Sep 28, 2023 · Sep 18, 2023 · Sep 18, 2023 · Sep 18, 2023
commit f208ccf2d9fb549c02a6fca788c95ad10f251aaa
diff --git a/docarray/utils/create_dynamic_doc_class.py b/docarray/utils/create_dynamic_doc_class.py
@@ -50,10 +50,6 @@ class MyDoc(BaseDoc):
     :param model: The input model
     :return: A new subclass of BaseDoc, where every DocList type in the schema is replaced by List.
     """
-    #    if is_pydantic_v2:
-    #        raise NotImplementedError(
-    #            'This method is not supported in Pydantic 2.0. Please use Pydantic 1.8.2 or lower.'
-    #        )
 
     fields: Dict[str, Any] = {}
     for field_name, field in model.__annotations__.items():
@@ -279,17 +275,42 @@ class MyDoc(BaseDoc):
             num_recursions=0,
             definitions=definitions,
         )
-        fields[field_name] = (
-            field_type,
-            FieldInfo(default=field_schema.pop('default', None), **field_schema),
-        )
+        if not is_pydantic_v2:
+            field_schema['default'] = field_schema.get('default', None)
+            fields[field_name] = (
+                field_type,
+                FieldInfo(**field_schema),
+            )
+        else:
+            field_kwargs = {}
+            field_json_schema_extra = {}
+            for k, v in field_schema.items():
+                if k in FieldInfo.__slots__:
+                    field_kwargs[k] = v
+                else:
+                    field_json_schema_extra[k] = v
+            fields[field_name] = (
+                field_type,
+                FieldInfo(
+                    json_schema_extra=field_json_schema_extra,
+                    **field_kwargs,
+                ),
+            )
 
     model = create_model(base_doc_name, __base__=BaseDoc, **fields)
-    # model.__config__.title = schema.get('title', model.__config__.title)
+    if not is_pydantic_v2:
+        model.__config__.title = schema.get('title', model.__config__.title)
+    else:
+        set_title = schema.get('title', model.model_config.get('title', None))
+        if set_title:
+            model.model_config['title'] = set_title
 
     for k in RESERVED_KEYS:
         if k in schema:
             schema.pop(k)
-    # model.__config__.schema_extra = schema
+    if not is_pydantic_v2:
+        model.__config__.schema_extra = schema
+    else:
+        model.model_config['json_schema_extra'] = schema
     cached_models[base_doc_name] = model
     return model
diff --git a/tests/units/util/test_create_dynamic_code_class.py b/tests/units/util/test_create_dynamic_code_class.py
@@ -7,11 +7,11 @@
 from docarray import BaseDoc, DocList
 from docarray.documents import TextDoc
 from docarray.typing import AnyTensor, ImageUrl
-from docarray.utils._internal.pydantic import is_pydantic_v2
 from docarray.utils.create_dynamic_doc_class import (
     create_base_doc_from_schema,
     create_pure_python_type_model,
 )
+from docarray.utils._internal.pydantic import is_pydantic_v2
 
 
 @pytest.mark.parametrize('transformation', ['proto', 'json'])
@@ -25,7 +25,7 @@ class Nested1Doc(BaseDoc):
         classvar: ClassVar[str] = 'classvar1'
 
     class CustomDoc(BaseDoc):
-        tensor: Optional[AnyTensor]
+        tensor: Optional[AnyTensor] = None
         url: ImageUrl
         lll: List[List[List[int]]] = [[[5]]]
         fff: List[List[List[float]]] = [[[5.2]]]
@@ -275,7 +275,8 @@ class CustomDoc(BaseDoc):
         new_custom_doc_model.schema().get('properties')['a']['another_extra']
         == 'I am another extra'
     )
-    assert (
-        new_custom_doc_model.schema().get('description')
-        == 'Here I have the description of the class'
-    )
+    if not is_pydantic_v2:
+        assert (
+            new_custom_doc_model.schema().get('description')
+            == 'Here I have the description of the class'
+        )