samsja · Sep 19, 2023 · Sep 19, 2023 · Sep 19, 2023 · Sep 19, 2023 · Sep 22, 2023
diff --git a/docarray/documents/base_predefined_doc.py b/docarray/documents/base_predefined_doc.py
@@ -0,0 +1,45 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+from docarray.base_doc.doc import BaseDoc
+from docarray.utils._internal.pydantic import is_pydantic_v2
+
+if is_pydantic_v2:
+    from pydantic_core import core_schema
+
+
+class PredefinedDoc(BaseDoc, ABC):
+    """
+    Custom class for handling predefined documents and can override their init and validation input.
+    """
+
+    @classmethod
+    @abstractmethod
+    def _docarray_custom_val(cls, value: Any) -> Dict[str, Any]:
+        ...
+
+    if is_pydantic_v2:
+
+        @classmethod
+        def __get_pydantic_core_schema__(
+            cls, source, handler
+        ) -> core_schema.CoreSchema:
+            if '__pydantic_core_schema__' in cls.__dict__:
+                if not cls.__pydantic_generic_metadata__['origin']:
+                    schema = cls.__pydantic_core_schema__
+
+            schema = handler(source)
+            return core_schema.general_wrap_validator_function(
+                function=cls._docarray_default_value_validation, schema=schema
+            )
+
+        @classmethod
+        def _docarray_default_value_validation(cls, value, model_validator, _):
+            # print(value)
+            return model_validator(cls._docarray_custom_val(value))
+
+    else:
+
+        @classmethod
+        def validate(cls, value: Any):
+            return super().validate(cls._docarray_custom_val(value))
diff --git a/docarray/documents/text.py b/docarray/documents/text.py
@@ -1,15 +1,19 @@
-from typing import Any, Optional, Type, TypeVar, Union
+from typing import Any, Dict, Optional, TypeVar
 
 from pydantic import Field
 
-from docarray.base_doc import BaseDoc
+from docarray.documents.base_predefined_doc import PredefinedDoc
 from docarray.typing import TextUrl
 from docarray.typing.tensor.embedding import AnyEmbedding
+from docarray.utils._internal.pydantic import is_pydantic_v2
+
+if is_pydantic_v2:
+    pass
 
 T = TypeVar('T', bound='TextDoc')
 
 
-class TextDoc(BaseDoc):
+class TextDoc(PredefinedDoc):
     """
     Document for handling text.
 
@@ -124,19 +128,27 @@ class MultiModalDoc(BaseDoc):
         default=None,
     )
 
-    def __init__(self, text: Optional[str] = None, **kwargs):
-        if 'text' not in kwargs:
-            kwargs['text'] = text
-        super().__init__(**kwargs)
+    def __init__(__pydantic_self__, *arg, **data) -> None:  # type: ignore
+        """Custom init function that allow validation without a dict as input but just a default value"""
+        # `__tracebackhide__` tells pytest and some other tools to omit this function from tracebacks
+
+        if len(arg) == 1 and not data:
+            data = {'text': arg[0]}
+        elif len(arg) > 1:
+            raise TypeError(
+                f'__init__() takes 1 positional argument but {len(arg)} were given'
+            )
+
+        super().__init__(**data)
 
     @classmethod
-    def validate(
-        cls: Type[T],
-        value: Union[str, Any],
-    ) -> T:
-        if isinstance(value, str):
-            value = cls(text=value)
-        return super().validate(value)
+    def _docarray_custom_val(cls, value: Any) -> Dict[str, Any]:
+        if not isinstance(value, dict):
+            new_value = {'text': value}
+        else:
+            new_value = value
+        print(new_value)
+        return new_value
 
     def __eq__(self, other: Any) -> bool:
         if isinstance(other, str):

diff --git a/tests/integrations/predefined_document/test_text.py b/tests/integrations/predefined_document/test_text.py
@@ -1,24 +1,19 @@
-import pytest
 from pydantic import parse_obj_as
 
 from docarray import BaseDoc
 from docarray.documents import TextDoc
-from docarray.utils._internal.pydantic import is_pydantic_v2
 
 
-@pytest.mark.skipif(is_pydantic_v2, reason="Not working with pydantic v2 for now")
 def test_simple_init():
     t = TextDoc(text='hello')
     assert t.text == 'hello'
 
 
-@pytest.mark.skipif(is_pydantic_v2, reason="Not working with pydantic v2 for now")
 def test_str_init():
     t = parse_obj_as(TextDoc, 'hello')
     assert t.text == 'hello'
 
 
-@pytest.mark.skipif(is_pydantic_v2, reason="Not working with pydantic v2 for now")
 def test_doc():
     class MyDoc(BaseDoc):
         text1: TextDoc
@@ -28,3 +23,14 @@ class MyDoc(BaseDoc):
 
     assert doc.text1.text == 'hello'
     assert doc.text2.text == 'world'
+
+
+def test_doc_2():
+    class MyDoc(BaseDoc):
+        text1: TextDoc
+        text2: TextDoc
+
+    doc = MyDoc(text1='hello1', text2='hello2')
+
+    assert doc.text1.text == 'hello1'
+    assert doc.text2.text == 'hello2'
diff --git a/tests/units/document/test_predefined_docs.py b/tests/units/document/test_predefined_docs.py
@@ -0,0 +1,25 @@
+from pydantic import parse_obj_as
+
+from docarray import BaseDoc
+from docarray.documents import TextDoc
+
+
+def test_nested_defaut_value():
+    class MyDocument(BaseDoc):
+        caption: TextDoc
+
+    doc = MyDocument(caption='A tiger in the jungle')
+
+    assert doc.caption.text == 'A tiger in the jungle'
+
+
+def test_parse_default_val():
+
+    doc = parse_obj_as(TextDoc, 'A tiger in the jungle')
+    assert doc.text == 'A tiger in the jungle'
+
+
+def test_parse_default_val_init():
+
+    doc = TextDoc('A tiger in the jungle')
+    assert doc.text == 'A tiger in the jungle'