From df3f45dfa6bd35f90fb776ab36daf357cc04d664 Mon Sep 17 00:00:00 2001 From: numb3r3 Date: Fri, 21 Jan 2022 14:50:39 +0800 Subject: [PATCH 1/3] fix(pydantic): disable smart union --- docarray/document/pydantic_model.py | 2 +- tests/unit/test_pydantic.py | 9 ++++++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/docarray/document/pydantic_model.py b/docarray/document/pydantic_model.py index b332c578a7c..c0e109cf413 100644 --- a/docarray/document/pydantic_model.py +++ b/docarray/document/pydantic_model.py @@ -38,7 +38,7 @@ class PydanticDocument(BaseModel): text: Optional[str] weight: Optional[float] uri: Optional[str] - tags: Optional[Dict[str, '_StructValueType']] + tags: Optional[Dict[str, Any]] offset: Optional[float] location: Optional[List[float]] embedding: Optional[Any] diff --git a/tests/unit/test_pydantic.py b/tests/unit/test_pydantic.py index 074f162852f..6255ebe117a 100644 --- a/tests/unit/test_pydantic.py +++ b/tests/unit/test_pydantic.py @@ -123,7 +123,14 @@ def test_with_embedding_no_tensor(): @pytest.mark.parametrize( 'tag_value, tag_type', - [(3, float), (3.4, float), ('hello', str), (True, bool), (False, bool)], + [ + (3.0, float), + (3.4, float), + ('hello', str), + ('1', str), + (True, bool), + (False, bool), + ], ) @pytest.mark.parametrize('protocol', ['protobuf', 'jsonschema']) def test_tags_int_float_str_bool(tag_type, tag_value, protocol): From 27fe344b9eed802cba9166fc0d7b99e15eaebb8d Mon Sep 17 00:00:00 2001 From: numb3r3 Date: Fri, 21 Jan 2022 14:53:53 +0800 Subject: [PATCH 2/3] fix: clean test case --- tests/unit/test_pydantic.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unit/test_pydantic.py b/tests/unit/test_pydantic.py index 6255ebe117a..174caafedad 100644 --- a/tests/unit/test_pydantic.py +++ b/tests/unit/test_pydantic.py @@ -125,7 +125,6 @@ def test_with_embedding_no_tensor(): 'tag_value, tag_type', [ (3.0, float), - (3.4, float), ('hello', str), ('1', str), (True, bool), From 359328ccdca54cd11bac213ef161598fa5d387a8 Mon Sep 17 00:00:00 2001 From: Han Xiao Date: Sun, 23 Jan 2022 10:31:21 +0100 Subject: [PATCH 3/3] fix(document): add smart union to pydantic document --- docarray/document/pydantic_model.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/docarray/document/pydantic_model.py b/docarray/document/pydantic_model.py index c0e109cf413..6661de63f2c 100644 --- a/docarray/document/pydantic_model.py +++ b/docarray/document/pydantic_model.py @@ -38,7 +38,7 @@ class PydanticDocument(BaseModel): text: Optional[str] weight: Optional[float] uri: Optional[str] - tags: Optional[Dict[str, Any]] + tags: Optional[Dict[str, '_StructValueType']] offset: Optional[float] location: Optional[List[float]] embedding: Optional[Any] @@ -59,6 +59,9 @@ def _blob2base64(cls, v): else: raise ValueError('must be bytes') + class Config: + smart_union = True + PydanticDocument.update_forward_refs()