fix(union_serializer): do not raise warnings in nested unions (#1513)

lukapeschke · web-flow · commit 2419981d9af7 · 2024-11-07T10:53:30.000Z
Signed-off-by: Luka Peschke &lt;mail@lukapeschke.com&gt;
diff --git a/src/serializers/type_serializers/union.rs b/src/serializers/type_serializers/union.rs
@@ -9,6 +9,7 @@ use crate::build_tools::py_schema_err;
 use crate::common::union::{Discriminator, SMALL_UNION_THRESHOLD};
 use crate::definitions::DefinitionsBuilder;
 use crate::tools::{truncate_safe_repr, SchemaDict};
+use crate::PydanticSerializationUnexpectedValue;
 
 use super::{
     infer_json_key, infer_serialize, infer_to_python, BuildSerializer, CombinedSerializer, Extra, SerCheck,
@@ -89,7 +90,8 @@ fn to_python(
         }
     }
 
-    if retry_with_lax_check {
+    // If extra.check is SerCheck::Strict, we're in a nested union
+    if extra.check != SerCheck::Strict && retry_with_lax_check {
         new_extra.check = SerCheck::Lax;
         for comb_serializer in choices {
             if let Ok(v) = comb_serializer.to_python(value, include, exclude, &new_extra) {
@@ -98,8 +100,17 @@ fn to_python(
         }
     }
 
-    for err in &errors {
-        extra.warnings.custom_warning(err.to_string());
+    // If extra.check is SerCheck::None, we're in a top-level union. We should thus raise the warnings
+    if extra.check == SerCheck::None {
+        for err in &errors {
+            extra.warnings.custom_warning(err.to_string());
+        }
+    }
+    // Otherwise, if we've encountered errors, return them to the parent union, which should take
+    // care of the formatting for us
+    else if !errors.is_empty() {
+        let message = errors.iter().map(ToString::to_string).collect::<Vec<_>>().join("\n");
+        return Err(PydanticSerializationUnexpectedValue::new_err(Some(message)));
     }
 
     infer_to_python(value, include, exclude, extra)
@@ -122,7 +133,8 @@ fn json_key<'a>(
         }
     }
 
-    if retry_with_lax_check {
+    // If extra.check is SerCheck::Strict, we're in a nested union
+    if extra.check != SerCheck::Strict && retry_with_lax_check {
         new_extra.check = SerCheck::Lax;
         for comb_serializer in choices {
             if let Ok(v) = comb_serializer.json_key(key, &new_extra) {
@@ -131,10 +143,18 @@ fn json_key<'a>(
         }
     }
 
-    for err in &errors {
-        extra.warnings.custom_warning(err.to_string());
+    // If extra.check is SerCheck::None, we're in a top-level union. We should thus raise the warnings
+    if extra.check == SerCheck::None {
+        for err in &errors {
+            extra.warnings.custom_warning(err.to_string());
+        }
+    }
+    // Otherwise, if we've encountered errors, return them to the parent union, which should take
+    // care of the formatting for us
+    else if !errors.is_empty() {
+        let message = errors.iter().map(ToString::to_string).collect::<Vec<_>>().join("\n");
+        return Err(PydanticSerializationUnexpectedValue::new_err(Some(message)));
     }
-
     infer_json_key(key, extra)
 }
 
@@ -160,7 +180,8 @@ fn serde_serialize<S: serde::ser::Serializer>(
         }
     }
 
-    if retry_with_lax_check {
+    // If extra.check is SerCheck::Strict, we're in a nested union
+    if extra.check != SerCheck::Strict && retry_with_lax_check {
         new_extra.check = SerCheck::Lax;
         for comb_serializer in choices {
             if let Ok(v) = comb_serializer.to_python(value, include, exclude, &new_extra) {
@@ -169,8 +190,14 @@ fn serde_serialize<S: serde::ser::Serializer>(
         }
     }
 
-    for err in &errors {
-        extra.warnings.custom_warning(err.to_string());
+    // If extra.check is SerCheck::None, we're in a top-level union. We should thus raise the warnings
+    if extra.check == SerCheck::None {
+        for err in &errors {
+            extra.warnings.custom_warning(err.to_string());
+        }
+    } else {
+        // NOTE: if this function becomes recursive at some point, an `Err(_)` containing the errors
+        // will have to be returned here
     }
 
     infer_serialize(value, serializer, include, exclude, extra)
diff --git a/tests/serializers/test_union.py b/tests/serializers/test_union.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import dataclasses
 import json
 import uuid
@@ -778,3 +780,223 @@ class ModelB:
     model_b = ModelB(field=1)
     assert s.to_python(model_a) == {'field': 1, 'TAG': 'a'}
     assert s.to_python(model_b) == {'field': 1, 'TAG': 'b'}
+
+
+class ModelDog:
+    def __init__(self, type_: Literal['dog']) -> None:
+        self.type_ = 'dog'
+
+
+class ModelCat:
+    def __init__(self, type_: Literal['cat']) -> None:
+        self.type_ = 'cat'
+
+
+class ModelAlien:
+    def __init__(self, type_: Literal['alien']) -> None:
+        self.type_ = 'alien'
+
+
+@pytest.fixture
+def model_a_b_union_schema() -> core_schema.UnionSchema:
+    return core_schema.union_schema(
+        [
+            core_schema.model_schema(
+                cls=ModelA,
+                schema=core_schema.model_fields_schema(
+                    fields={
+                        'a': core_schema.model_field(core_schema.str_schema()),
+                        'b': core_schema.model_field(core_schema.str_schema()),
+                    },
+                ),
+            ),
+            core_schema.model_schema(
+                cls=ModelB,
+                schema=core_schema.model_fields_schema(
+                    fields={
+                        'c': core_schema.model_field(core_schema.str_schema()),
+                        'd': core_schema.model_field(core_schema.str_schema()),
+                    },
+                ),
+            ),
+        ]
+    )
+
+
+@pytest.fixture
+def union_of_unions_schema(model_a_b_union_schema: core_schema.UnionSchema) -> core_schema.UnionSchema:
+    return core_schema.union_schema(
+        [
+            model_a_b_union_schema,
+            core_schema.union_schema(
+                [
+                    core_schema.model_schema(
+                        cls=ModelCat,
+                        schema=core_schema.model_fields_schema(
+                            fields={
+                                'type_': core_schema.model_field(core_schema.literal_schema(['cat'])),
+                            },
+                        ),
+                    ),
+                    core_schema.model_schema(
+                        cls=ModelDog,
+                        schema=core_schema.model_fields_schema(
+                            fields={
+                                'type_': core_schema.model_field(core_schema.literal_schema(['dog'])),
+                            },
+                        ),
+                    ),
+                ]
+            ),
+        ]
+    )
+
+
+@pytest.mark.parametrize(
+    'input,expected',
+    [
+        (ModelA(a='a', b='b'), {'a': 'a', 'b': 'b'}),
+        (ModelB(c='c', d='d'), {'c': 'c', 'd': 'd'}),
+        (ModelCat(type_='cat'), {'type_': 'cat'}),
+        (ModelDog(type_='dog'), {'type_': 'dog'}),
+    ],
+)
+def test_union_of_unions_of_models(union_of_unions_schema: core_schema.UnionSchema, input: Any, expected: Any) -> None:
+    s = SchemaSerializer(union_of_unions_schema)
+    assert s.to_python(input, warnings='error') == expected
+
+
+def test_union_of_unions_of_models_invalid_variant(union_of_unions_schema: core_schema.UnionSchema) -> None:
+    s = SchemaSerializer(union_of_unions_schema)
+    # All warnings should be available
+    messages = [
+        'Expected `ModelA` but got `ModelAlien`',
+        'Expected `ModelB` but got `ModelAlien`',
+        'Expected `ModelCat` but got `ModelAlien`',
+        'Expected `ModelDog` but got `ModelAlien`',
+    ]
+
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter('always')
+        s.to_python(ModelAlien(type_='alien'))
+        for m in messages:
+            assert m in str(w[0].message)
+
+
+@pytest.fixture
+def tagged_union_of_unions_schema(model_a_b_union_schema: core_schema.UnionSchema) -> core_schema.UnionSchema:
+    return core_schema.union_schema(
+        [
+            model_a_b_union_schema,
+            core_schema.tagged_union_schema(
+                discriminator='type_',
+                choices={
+                    'cat': core_schema.model_schema(
+                        cls=ModelCat,
+                        schema=core_schema.model_fields_schema(
+                            fields={
+                                'type_': core_schema.model_field(core_schema.literal_schema(['cat'])),
+                            },
+                        ),
+                    ),
+                    'dog': core_schema.model_schema(
+                        cls=ModelDog,
+                        schema=core_schema.model_fields_schema(
+                            fields={
+                                'type_': core_schema.model_field(core_schema.literal_schema(['dog'])),
+                            },
+                        ),
+                    ),
+                },
+            ),
+        ]
+    )
+
+
+@pytest.mark.parametrize(
+    'input,expected',
+    [
+        (ModelA(a='a', b='b'), {'a': 'a', 'b': 'b'}),
+        (ModelB(c='c', d='d'), {'c': 'c', 'd': 'd'}),
+        (ModelCat(type_='cat'), {'type_': 'cat'}),
+        (ModelDog(type_='dog'), {'type_': 'dog'}),
+    ],
+)
+def test_union_of_unions_of_models_with_tagged_union(
+    tagged_union_of_unions_schema: core_schema.UnionSchema, input: Any, expected: Any
+) -> None:
+    s = SchemaSerializer(tagged_union_of_unions_schema)
+    assert s.to_python(input, warnings='error') == expected
+
+
+def test_union_of_unions_of_models_with_tagged_union_invalid_variant(
+    tagged_union_of_unions_schema: core_schema.UnionSchema,
+) -> None:
+    s = SchemaSerializer(tagged_union_of_unions_schema)
+    # All warnings should be available
+    messages = [
+        'Expected `ModelA` but got `ModelAlien`',
+        'Expected `ModelB` but got `ModelAlien`',
+        'Expected `ModelCat` but got `ModelAlien`',
+        'Expected `ModelDog` but got `ModelAlien`',
+    ]
+
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter('always')
+        s.to_python(ModelAlien(type_='alien'))
+        for m in messages:
+            assert m in str(w[0].message)
+
+
+@pytest.mark.parametrize(
+    'input,expected',
+    [
+        ({True: '1'}, b'{"true":"1"}'),
+        ({1: '1'}, b'{"1":"1"}'),
+        ({2.3: '1'}, b'{"2.3":"1"}'),
+        ({'a': 'b'}, b'{"a":"b"}'),
+    ],
+)
+def test_union_of_unions_of_models_with_tagged_union_json_key_serialization(
+    input: dict[bool | int | float | str, str], expected: bytes
+) -> None:
+    s = SchemaSerializer(
+        core_schema.dict_schema(
+            keys_schema=core_schema.union_schema(
+                [
+                    core_schema.union_schema([core_schema.bool_schema(), core_schema.int_schema()]),
+                    core_schema.union_schema([core_schema.float_schema(), core_schema.str_schema()]),
+                ]
+            ),
+            values_schema=core_schema.str_schema(),
+        )
+    )
+
+    assert s.to_json(input, warnings='error') == expected
+
+
+@pytest.mark.parametrize(
+    'input,expected',
+    [
+        ({'key': True}, b'{"key":true}'),
+        ({'key': 1}, b'{"key":1}'),
+        ({'key': 2.3}, b'{"key":2.3}'),
+        ({'key': 'a'}, b'{"key":"a"}'),
+    ],
+)
+def test_union_of_unions_of_models_with_tagged_union_json_serialization(
+    input: dict[str, bool | int | float | str], expected: bytes
+) -> None:
+    s = SchemaSerializer(
+        core_schema.dict_schema(
+            keys_schema=core_schema.str_schema(),
+            values_schema=core_schema.union_schema(
+                [
+                    core_schema.union_schema([core_schema.bool_schema(), core_schema.int_schema()]),
+                    core_schema.union_schema([core_schema.float_schema(), core_schema.str_schema()]),
+                ]
+            ),
+        )
+    )
+
+    assert s.to_json(input, warnings='error') == expected

Original file line number	Diff line number	Diff line change
`@@ -9,6 +9,7 @@ use crate::build_tools::py_schema_err;`
`9`	`9`	`use crate::common::union::{Discriminator, SMALL_UNION_THRESHOLD};`
`10`	`10`	`use crate::definitions::DefinitionsBuilder;`
`11`	`11`	`use crate::tools::{truncate_safe_repr, SchemaDict};`
	`12`	`+use crate::PydanticSerializationUnexpectedValue;`
`12`	`13`
`13`	`14`	`use super::{`
`14`	`15`	`infer_json_key, infer_serialize, infer_to_python, BuildSerializer, CombinedSerializer, Extra, SerCheck,`
`@@ -89,7 +90,8 @@ fn to_python(`
`89`	`90`	`}`
`90`	`91`	`}`
`91`	`92`
`92`		`- if retry_with_lax_check {`
	`93`	`+ // If extra.check is SerCheck::Strict, we're in a nested union`
	`94`	`+ if extra.check != SerCheck::Strict && retry_with_lax_check {`
`93`	`95`	`new_extra.check = SerCheck::Lax;`
`94`	`96`	`for comb_serializer in choices {`
`95`	`97`	`if let Ok(v) = comb_serializer.to_python(value, include, exclude, &new_extra) {`
`@@ -98,8 +100,17 @@ fn to_python(`
`98`	`100`	`}`
`99`	`101`	`}`
`100`	`102`
`101`		`- for err in &errors {`
`102`		`- extra.warnings.custom_warning(err.to_string());`
	`103`	`+ // If extra.check is SerCheck::None, we're in a top-level union. We should thus raise the warnings`
	`104`	`+ if extra.check == SerCheck::None {`
	`105`	`+ for err in &errors {`
	`106`	`+ extra.warnings.custom_warning(err.to_string());`
	`107`	`+ }`
	`108`	`+ }`
	`109`	`+ // Otherwise, if we've encountered errors, return them to the parent union, which should take`
	`110`	`+ // care of the formatting for us`
	`111`	`+ else if !errors.is_empty() {`
	`112`	`+ let message = errors.iter().map(ToString::to_string).collect::<Vec<_>>().join("\n");`
	`113`	`+ return Err(PydanticSerializationUnexpectedValue::new_err(Some(message)));`
`103`	`114`	`}`
`104`	`115`
`105`	`116`	`infer_to_python(value, include, exclude, extra)`
`@@ -122,7 +133,8 @@ fn json_key<'a>(`
`122`	`133`	`}`
`123`	`134`	`}`
`124`	`135`
`125`		`- if retry_with_lax_check {`
	`136`	`+ // If extra.check is SerCheck::Strict, we're in a nested union`
	`137`	`+ if extra.check != SerCheck::Strict && retry_with_lax_check {`
`126`	`138`	`new_extra.check = SerCheck::Lax;`
`127`	`139`	`for comb_serializer in choices {`
`128`	`140`	`if let Ok(v) = comb_serializer.json_key(key, &new_extra) {`
`@@ -131,10 +143,18 @@ fn json_key<'a>(`
`131`	`143`	`}`
`132`	`144`	`}`
`133`	`145`
`134`		`- for err in &errors {`
`135`		`- extra.warnings.custom_warning(err.to_string());`
	`146`	`+ // If extra.check is SerCheck::None, we're in a top-level union. We should thus raise the warnings`
	`147`	`+ if extra.check == SerCheck::None {`
	`148`	`+ for err in &errors {`
	`149`	`+ extra.warnings.custom_warning(err.to_string());`
	`150`	`+ }`
	`151`	`+ }`
	`152`	`+ // Otherwise, if we've encountered errors, return them to the parent union, which should take`
	`153`	`+ // care of the formatting for us`
	`154`	`+ else if !errors.is_empty() {`
	`155`	`+ let message = errors.iter().map(ToString::to_string).collect::<Vec<_>>().join("\n");`
	`156`	`+ return Err(PydanticSerializationUnexpectedValue::new_err(Some(message)));`
`136`	`157`	`}`
`137`		`-`
`138`	`158`	`infer_json_key(key, extra)`
`139`	`159`	`}`
`140`	`160`
`@@ -160,7 +180,8 @@ fn serde_serialize<S: serde::ser::Serializer>(`
`160`	`180`	`}`
`161`	`181`	`}`
`162`	`182`
`163`		`- if retry_with_lax_check {`
	`183`	`+ // If extra.check is SerCheck::Strict, we're in a nested union`
	`184`	`+ if extra.check != SerCheck::Strict && retry_with_lax_check {`
`164`	`185`	`new_extra.check = SerCheck::Lax;`
`165`	`186`	`for comb_serializer in choices {`
`166`	`187`	`if let Ok(v) = comb_serializer.to_python(value, include, exclude, &new_extra) {`
`@@ -169,8 +190,14 @@ fn serde_serialize<S: serde::ser::Serializer>(`
`169`	`190`	`}`
`170`	`191`	`}`
`171`	`192`
`172`		`- for err in &errors {`
`173`		`- extra.warnings.custom_warning(err.to_string());`
	`193`	`+ // If extra.check is SerCheck::None, we're in a top-level union. We should thus raise the warnings`
	`194`	`+ if extra.check == SerCheck::None {`
	`195`	`+ for err in &errors {`
	`196`	`+ extra.warnings.custom_warning(err.to_string());`
	`197`	`+ }`
	`198`	`+ } else {`
	`199`	+ // NOTE: if this function becomes recursive at some point, an `Err(_)` containing the errors
	`200`	`+ // will have to be returned here`
`174`	`201`	`}`
`175`	`202`
`176`	`203`	`infer_serialize(value, serializer, include, exclude, extra)`