custom errors on union and tagged-union (#262)

* custom errors on union and tagged-union

* use PydanticValueError directly

* formatting tweak

Author: samuelcolvin

pydantic_core/core_schema.py

@@ -527,9 +527,15 @@ def nullable_schema(schema: CoreSchema, *, strict: bool | None = None, ref: str
     return dict_not_none(type='nullable', schema=schema, strict=strict, ref=ref)
 
 
+class CustomError(TypedDict):
+    kind: str
+    message: str
+
+
 class UnionSchema(TypedDict, total=False):
     type: Required[Literal['union']]
     choices: Required[List[CoreSchema]]
+    custom_error: CustomError
     strict: bool
     ref: str
 
@@ -542,6 +548,7 @@ class TaggedUnionSchema(TypedDict):
     type: Literal['tagged-union']
     choices: Dict[str, CoreSchema]
     discriminator: Union[str, List[Union[str, int]], List[List[Union[str, int]]], Callable[[Any], Optional[str]]]
+    custom_error: NotRequired[CustomError]
     strict: NotRequired[bool]
     ref: NotRequired[str]
 

src/errors/value_exception.rs

@@ -7,8 +7,7 @@ use crate::input::Input;
 use super::{ErrorKind, ValError};
 
 #[pyclass(extends=PyValueError, module="pydantic_core._pydantic_core")]
-#[derive(Clone)]
-#[cfg_attr(debug_assertions, derive(Debug))]
+#[derive(Debug, Clone)]
 pub struct PydanticValueError {
     kind: String,
     message_template: String,
@@ -18,7 +17,7 @@ pub struct PydanticValueError {
 #[pymethods]
 impl PydanticValueError {
     #[new]
-    fn py_new(py: Python, kind: String, message_template: String, context: Option<&PyDict>) -> Self {
+    pub fn py_new(py: Python, kind: String, message_template: String, context: Option<&PyDict>) -> Self {
         Self {
             kind,
             message_template,

src/validators/union.rs

@@ -8,7 +8,7 @@ use pyo3::types::{PyDict, PyList, PyString};
 use ahash::AHashMap;
 
 use crate::build_tools::{is_strict, schema_or_config, SchemaDict};
-use crate::errors::{ErrorKind, ValError, ValLineError, ValResult};
+use crate::errors::{ErrorKind, PydanticValueError, ValError, ValLineError, ValResult};
 use crate::input::{GenericMapping, Input};
 use crate::lookup_key::LookupKey;
 use crate::questions::Question;
@@ -19,6 +19,7 @@ use super::{build_validator, BuildContext, BuildValidator, CombinedValidator, Ex
 #[derive(Debug, Clone)]
 pub struct UnionValidator {
     choices: Vec<CombinedValidator>,
+    custom_error: Option<PydanticValueError>,
     strict: bool,
     name: String,
 }
@@ -31,8 +32,9 @@ impl BuildValidator for UnionValidator {
         config: Option<&PyDict>,
         build_context: &mut BuildContext,
     ) -> PyResult<CombinedValidator> {
+        let py = schema.py();
         let choices: Vec<CombinedValidator> = schema
-            .get_as_req::<&PyList>(intern!(schema.py(), "choices"))?
+            .get_as_req::<&PyList>(intern!(py, "choices"))?
             .iter()
             .map(|choice| build_validator(choice, config, build_context))
             .collect::<PyResult<Vec<CombinedValidator>>>()?;
@@ -41,13 +43,41 @@ impl BuildValidator for UnionValidator {
 
         Ok(Self {
             choices,
+            custom_error: get_custom_error(py, schema)?,
             strict: is_strict(schema, config)?,
             name: format!("{}[{}]", Self::EXPECTED_TYPE, descr),
         }
         .into())
     }
 }
 
+fn get_custom_error(py: Python, schema: &PyDict) -> PyResult<Option<PydanticValueError>> {
+    match schema.get_as::<&PyDict>(intern!(py, "custom_error"))? {
+        Some(custom_error) => Ok(Some(PydanticValueError::py_new(
+            py,
+            custom_error.get_as_req::<String>(intern!(py, "kind"))?,
+            custom_error.get_as_req::<String>(intern!(py, "message"))?,
+            None,
+        ))),
+        None => Ok(None),
+    }
+}
+
+impl UnionValidator {
+    fn or_custom_error<'s, 'data>(
+        &'s self,
+        errors: Option<Vec<ValLineError<'data>>>,
+        input: &'data impl Input<'data>,
+    ) -> ValError<'data> {
+        if let Some(errors) = errors {
+            ValError::LineErrors(errors)
+        } else {
+            let value_error = self.custom_error.as_ref().unwrap();
+            value_error.clone().into_val_error(input)
+        }
+    }
+}
+
 impl Validator for UnionValidator {
     fn validate<'s, 'data>(
         &'s self,
@@ -58,7 +88,10 @@ impl Validator for UnionValidator {
         recursion_guard: &'s mut RecursionGuard,
     ) -> ValResult<'data, PyObject> {
         if extra.strict.unwrap_or(self.strict) {
-            let mut errors: Vec<ValLineError> = Vec::with_capacity(self.choices.len());
+            let mut errors: Option<Vec<ValLineError>> = match self.custom_error {
+                Some(_) => None,
+                None => Some(Vec::with_capacity(self.choices.len())),
+            };
             let strict_extra = extra.as_strict();
 
             for validator in &self.choices {
@@ -67,14 +100,16 @@ impl Validator for UnionValidator {
                     otherwise => return otherwise,
                 };
 
-                errors.extend(
-                    line_errors
-                        .into_iter()
-                        .map(|err| err.with_outer_location(validator.get_name().into())),
-                );
+                if let Some(ref mut errors) = errors {
+                    errors.extend(
+                        line_errors
+                            .into_iter()
+                            .map(|err| err.with_outer_location(validator.get_name().into())),
+                    );
+                }
             }
 
-            Err(ValError::LineErrors(errors))
+            Err(self.or_custom_error(errors, input))
         } else {
             // 1st pass: check if the value is an exact instance of one of the Union types,
             // e.g. use validate in strict mode
@@ -88,7 +123,10 @@ impl Validator for UnionValidator {
                 return res;
             }
 
-            let mut errors: Vec<ValLineError> = Vec::with_capacity(self.choices.len());
+            let mut errors: Option<Vec<ValLineError>> = match self.custom_error {
+                Some(_) => None,
+                None => Some(Vec::with_capacity(self.choices.len())),
+            };
 
             // 2nd pass: check if the value can be coerced into one of the Union types, e.g. use validate
             for validator in &self.choices {
@@ -97,14 +135,16 @@ impl Validator for UnionValidator {
                     success => return success,
                 };
 
-                errors.extend(
-                    line_errors
-                        .into_iter()
-                        .map(|err| err.with_outer_location(validator.get_name().into())),
-                );
+                if let Some(ref mut errors) = errors {
+                    errors.extend(
+                        line_errors
+                            .into_iter()
+                            .map(|err| err.with_outer_location(validator.get_name().into())),
+                    );
+                }
             }
 
-            Err(ValError::LineErrors(errors))
+            Err(self.or_custom_error(errors, input))
         }
     }
 
@@ -160,6 +200,7 @@ pub struct TaggedUnionValidator {
     discriminator: Discriminator,
     from_attributes: bool,
     strict: bool,
+    custom_error: Option<PydanticValueError>,
     tags_repr: String,
     discriminator_repr: String,
     name: String,
@@ -206,6 +247,7 @@ impl BuildValidator for TaggedUnionValidator {
             discriminator,
             from_attributes,
             strict: is_strict(schema, config)?,
+            custom_error: get_custom_error(py, schema)?,
             tags_repr,
             discriminator_repr,
             name: format!("{}[{}]", Self::EXPECTED_TYPE, descr),
@@ -341,6 +383,8 @@ impl TaggedUnionValidator {
                 Ok(res) => Ok(res),
                 Err(err) => Err(err.with_outer_location(tag.as_ref().into())),
             }
+        } else if let Some(ref custom_error) = self.custom_error {
+            Err(custom_error.clone().into_val_error(input))
         } else {
             Err(ValError::new(
                 ErrorKind::UnionTagInvalid {
@@ -354,11 +398,15 @@ impl TaggedUnionValidator {
     }
 
     fn tag_not_found<'s, 'data>(&'s self, input: &'data impl Input<'data>) -> ValError<'data> {
-        ValError::new(
-            ErrorKind::UnionTagNotFound {
-                discriminator: self.discriminator_repr.clone(),
-            },
-            input,
-        )
+        if let Some(ref custom_error) = self.custom_error {
+            custom_error.clone().into_val_error(input)
+        } else {
+            ValError::new(
+                ErrorKind::UnionTagNotFound {
+                    discriminator: self.discriminator_repr.clone(),
+                },
+                input,
+            )
+        }
     }
 }

tests/validators/test_tagged_union.py

@@ -287,3 +287,48 @@ def test_downcast_error():
     v = SchemaValidator({'type': 'tagged-union', 'discriminator': lambda x: 123, 'choices': {'str': {'type': 'str'}}})
     with pytest.raises(TypeError, match="'int' object cannot be converted to 'PyString'"):
         v.validate_python('x')
+
+
+def test_custom_error():
+    v = SchemaValidator(
+        {
+            'type': 'tagged-union',
+            'discriminator': 'foo',
+            'custom_error': {'kind': 'snap', 'message': 'Input should be a foo or bar'},
+            'choices': {
+                'apple': {
+                    'type': 'typed-dict',
+                    'fields': {'foo': {'schema': {'type': 'str'}}, 'bar': {'schema': {'type': 'int'}}},
+                },
+                'banana': {
+                    'type': 'typed-dict',
+                    'fields': {
+                        'foo': {'schema': {'type': 'str'}},
+                        'spam': {'schema': {'type': 'list', 'items_schema': {'type': 'int'}}},
+                    },
+                },
+            },
+        }
+    )
+    assert v.validate_python({'foo': 'apple', 'bar': '123'}) == {'foo': 'apple', 'bar': 123}
+    with pytest.raises(ValidationError) as exc_info:
+        v.validate_python({'spam': 'apple', 'bar': 'Bar'})
+    # insert_assert(exc_info.value.errors())
+    assert exc_info.value.errors() == [
+        {
+            'kind': 'snap',
+            'loc': [],
+            'message': 'Input should be a foo or bar',
+            'input_value': {'spam': 'apple', 'bar': 'Bar'},
+        }
+    ]
+    with pytest.raises(ValidationError) as exc_info:
+        v.validate_python({'foo': 'other', 'bar': 'Bar'})
+    assert exc_info.value.errors() == [
+        {
+            'kind': 'snap',
+            'loc': [],
+            'message': 'Input should be a foo or bar',
+            'input_value': {'foo': 'other', 'bar': 'Bar'},
+        }
+    ]

tests/validators/test_union.py

@@ -244,3 +244,21 @@ def test_strict_union():
         {'kind': 'bool_type', 'loc': ['bool'], 'message': 'Input should be a valid boolean', 'input_value': '123'},
         {'kind': 'int_type', 'loc': ['int'], 'message': 'Input should be a valid integer', 'input_value': '123'},
     ]
+
+
+def test_custom_error():
+    v = SchemaValidator(
+        {
+            'type': 'union',
+            'choices': [{'type': 'str'}, {'type': 'bytes'}],
+            'custom_error': {'kind': 'my_error', 'message': 'Input should be a string or bytes'},
+        }
+    )
+    assert v.validate_python('hello') == 'hello'
+    assert v.validate_python(b'hello') == b'hello'
+    with pytest.raises(ValidationError) as exc_info:
+        v.validate_python(123)
+    # insert_assert(exc_info.value.errors())
+    assert exc_info.value.errors() == [
+        {'kind': 'my_error', 'loc': [], 'message': 'Input should be a string or bytes', 'input_value': 123}
+    ]