diff --git a/setup.py b/setup.py index 63ce1795748..c1fc9e3662f 100644 --- a/setup.py +++ b/setup.py @@ -48,6 +48,7 @@ 'matplotlib', 'Pillow', 'fastapi', + 'pydantic>=1.9.0', 'uvicorn', ], 'full': [ @@ -63,6 +64,7 @@ 'scipy', 'av', 'fastapi', + 'pydantic>=1.9.0', 'uvicorn', 'strawberry-graphql', ], diff --git a/tests/unit/test_pydantic.py b/tests/unit/test_pydantic.py index 9ff6c8b052a..da67954efe5 100644 --- a/tests/unit/test_pydantic.py +++ b/tests/unit/test_pydantic.py @@ -159,6 +159,22 @@ def test_tags_int_float_str_bool(tag_type, tag_value, protocol): assert isinstance(dd, tag_type) +@pytest.mark.parametrize('protocol', ['protobuf', 'jsonschema']) +def test_infinity_no_coercion(protocol): + # Test for issue #948: https://github.com/docarray/docarray/issues/948 + d = Document() + d.tags['title'] = 'Infinity' + + d_pydantic = d.to_pydantic_model() + d_pydantic.tags['title'] = 'Infinity' + + d_json = d.to_json(protocol=protocol) + assert '"title": "Infinity"' in d_json + + d_dict = d.to_dict(protocol=protocol) + assert d_dict['tags']['title'] == 'Infinity' + + @pytest.mark.parametrize( 'blob', [None, b'123', bytes(Document()), bytes(bytearray(os.urandom(512 * 4)))] )