diff --git a/docarray/documents/base_predefined_doc.py b/docarray/documents/base_predefined_doc.py new file mode 100644 index 0000000000..84c8919f70 --- /dev/null +++ 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 index 4df291bb18..dcd95d6221 100644 --- 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 index 5b89844ca3..5906217545 100644 --- 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 new file mode 100644 index 0000000000..14b429c774 --- /dev/null +++ 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'