Use ChoiceKey in TaggedUnionValidator for int keys (#405)

* Mostly working ChoiceKey for int keys

* Update src/validators/union.rs

Co-authored-by: Samuel Colvin <[email protected]>

* Update src/validators/union.rs

Co-authored-by: Samuel Colvin <[email protected]>

* Update src/validators/union.rs

Co-authored-by: Samuel Colvin <[email protected]>

* address comments

* address more comments

* Address more comments and fix test

* add tests for int choice keys

* fix

* add enum choices test

* update repeated tag test

* Set from_attributes to true by default for TaggedUnionValidator

* Get test passing

* fix tags_repr for ints

* refactor repeated tags test

* Use _extra.strict for literal validation of strs and ints

* support i64 location keys, tests

* change literal strict usage

---------

Co-authored-by: Samuel Colvin <[email protected]>
Co-authored-by: David Montague <[email protected]>

Author: 3 people

generate_self_schema.py

@@ -129,11 +129,14 @@ def type_dict_schema(typed_dict) -> dict[str, Any]:  # noqa: C901
                     schema = {'type': 'list', 'items_schema': schema_ref_validator}
                 elif fr_arg == 'Dict[str, CoreSchema]':
                     schema = {'type': 'dict', 'keys_schema': {'type': 'str'}, 'values_schema': schema_ref_validator}
-                elif fr_arg == 'Dict[str, Union[str, CoreSchema]]':
+                elif fr_arg == 'Dict[Union[str, int], Union[str, int, CoreSchema]]':
                     schema = {
                         'type': 'dict',
-                        'keys_schema': {'type': 'str'},
-                        'values_schema': {'type': 'union', 'choices': [{'type': 'str'}, schema_ref_validator]},
+                        'keys_schema': {'type': 'union', 'choices': [{'type': 'str'}, {'type': 'int'}]},
+                        'values_schema': {
+                            'type': 'union',
+                            'choices': [{'type': 'str'}, {'type': 'int'}, schema_ref_validator],
+                        },
                     }
                 else:
                     raise ValueError(f'Unknown Schema forward ref: {fr_arg}')

pydantic_core/core_schema.py

@@ -1805,27 +1805,29 @@ def union_schema(
 
 class TaggedUnionSchema(TypedDict, total=False):
     type: Required[Literal['tagged-union']]
-    choices: Required[Dict[str, Union[str, CoreSchema]]]
+    choices: Required[Dict[Union[str, int], Union[str, int, CoreSchema]]]
     discriminator: Required[
         Union[str, List[Union[str, int]], List[List[Union[str, int]]], Callable[[Any], Optional[str]]]
     ]
     custom_error_type: str
     custom_error_message: str
     custom_error_context: Dict[str, Union[str, int, float]]
     strict: bool
+    from_attributes: bool  # default: True
     ref: str
     metadata: Any
     serialization: SerSchema
 
 
 def tagged_union_schema(
-    choices: Dict[str, str | CoreSchema],
+    choices: Dict[Union[int, str], int | str | CoreSchema],
     discriminator: str | list[str | int] | list[list[str | int]] | Callable[[Any], str | None],
     *,
     custom_error_type: str | None = None,
     custom_error_message: str | None = None,
     custom_error_context: dict[str, int | str | float] | None = None,
     strict: bool | None = None,
+    from_attributes: bool | None = None,
     ref: str | None = None,
     metadata: Any = None,
     serialization: SerSchema | None = None,
@@ -1873,6 +1875,7 @@ def tagged_union_schema(
         custom_error_message: The custom error message to use if the validation fails
         custom_error_context: The custom error context to use if the validation fails
         strict: Whether the underlying schemas should be validated with strict mode
+        from_attributes: Whether to use the attributes of the object to retrieve the discriminator value
         ref: See [TODO] for details
         metadata: See [TODO] for details
         serialization: Custom serialization schema
@@ -1885,6 +1888,7 @@ def tagged_union_schema(
         custom_error_message=custom_error_message,
         custom_error_context=custom_error_context,
         strict=strict,
+        from_attributes=from_attributes,
         ref=ref,
         metadata=metadata,
         serialization=serialization,

src/errors/location.rs

@@ -11,8 +11,12 @@ use pyo3::types::PyTuple;
 pub enum LocItem {
     /// string type key, used to identify items from a dict or anything that implements `__getitem__`
     S(String),
-    /// integer key, used to get items from a list, tuple OR a dict with int keys `Dict[int, ...]` (python only)
-    I(usize),
+    /// integer key, used to get:
+    ///   * items from a list
+    ///   * items from a tuple
+    ///   * dict with int keys `Dict[int, ...]` (python only)
+    ///   * with integer keys in tagged unions
+    I(i64),
 }
 
 impl fmt::Display for LocItem {
@@ -36,12 +40,18 @@ impl From<&str> for LocItem {
     }
 }
 
-impl From<usize> for LocItem {
-    fn from(i: usize) -> Self {
+impl From<i64> for LocItem {
+    fn from(i: i64) -> Self {
         Self::I(i)
     }
 }
 
+impl From<usize> for LocItem {
+    fn from(u: usize) -> Self {
+        Self::I(u as i64)
+    }
+}
+
 impl ToPyObject for LocItem {
     fn to_object(&self, py: Python<'_>) -> PyObject {
         match self {

src/input/input_json.rs

@@ -23,7 +23,7 @@ impl<'a> Input<'a> for JsonInput {
     #[cfg_attr(has_no_coverage, no_coverage)]
     fn as_loc_item(&self) -> LocItem {
         match self {
-            JsonInput::Int(i) => LocItem::I(*i as usize),
+            JsonInput::Int(i) => (*i).into(),
             JsonInput::String(s) => s.as_str().into(),
             v => format!("{v:?}").into(),
         }

src/validators/literal.rs

@@ -71,11 +71,11 @@ impl Validator for LiteralSingleStringValidator {
         &'s self,
         py: Python<'data>,
         input: &'data impl Input<'data>,
-        _extra: &Extra,
+        extra: &Extra,
         _slots: &'data [CombinedValidator],
         _recursion_guard: &'s mut RecursionGuard,
     ) -> ValResult<'data, PyObject> {
-        let either_str = input.strict_str()?;
+        let either_str = input.validate_str(extra.strict.unwrap_or(false))?;
         if either_str.as_cow()?.as_ref() == self.expected.as_str() {
             Ok(input.to_object(py))
         } else {
@@ -113,12 +113,12 @@ impl Validator for LiteralSingleIntValidator {
         &'s self,
         py: Python<'data>,
         input: &'data impl Input<'data>,
-        _extra: &Extra,
+        extra: &Extra,
         _slots: &'data [CombinedValidator],
         _recursion_guard: &'s mut RecursionGuard,
     ) -> ValResult<'data, PyObject> {
-        let str = input.strict_int()?;
-        if str == self.expected {
+        let int = input.validate_int(extra.strict.unwrap_or(false))?;
+        if int == self.expected {
             Ok(input.to_object(py))
         } else {
             Err(ValError::new(
@@ -168,11 +168,11 @@ impl Validator for LiteralMultipleStringsValidator {
         &'s self,
         py: Python<'data>,
         input: &'data impl Input<'data>,
-        _extra: &Extra,
+        extra: &Extra,
         _slots: &'data [CombinedValidator],
         _recursion_guard: &'s mut RecursionGuard,
     ) -> ValResult<'data, PyObject> {
-        let either_str = input.strict_str()?;
+        let either_str = input.validate_str(extra.strict.unwrap_or(false))?;
         if self.expected.contains(either_str.as_cow()?.as_ref()) {
             Ok(input.to_object(py))
         } else {
@@ -223,11 +223,11 @@ impl Validator for LiteralMultipleIntsValidator {
         &'s self,
         py: Python<'data>,
         input: &'data impl Input<'data>,
-        _extra: &Extra,
+        extra: &Extra,
         _slots: &'data [CombinedValidator],
         _recursion_guard: &'s mut RecursionGuard,
     ) -> ValResult<'data, PyObject> {
-        let int = input.strict_int()?;
+        let int = input.validate_int(extra.strict.unwrap_or(false))?;
         if self.expected.contains(&int) {
             Ok(input.to_object(py))
         } else {
@@ -287,19 +287,20 @@ impl Validator for LiteralGeneralValidator {
         &'s self,
         py: Python<'data>,
         input: &'data impl Input<'data>,
-        _extra: &Extra,
+        extra: &Extra,
         _slots: &'data [CombinedValidator],
         _recursion_guard: &'s mut RecursionGuard,
     ) -> ValResult<'data, PyObject> {
+        let strict = extra.strict.unwrap_or(false);
         if !self.expected_int.is_empty() {
-            if let Ok(int) = input.strict_int() {
+            if let Ok(int) = input.validate_int(strict) {
                 if self.expected_int.contains(&int) {
                     return Ok(input.to_object(py));
                 }
             }
         }
         if !self.expected_str.is_empty() {
-            if let Ok(either_str) = input.strict_str() {
+            if let Ok(either_str) = input.validate_str(strict) {
                 if self.expected_str.contains(either_str.as_cow()?.as_ref()) {
                     return Ok(input.to_object(py));
                 }

src/validators/union.rs

@@ -1,14 +1,16 @@
 use std::borrow::Cow;
+use std::fmt;
 use std::fmt::Write;
 
+use pyo3::exceptions::PyTypeError;
 use pyo3::intern;
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyList, PyString};
 
 use ahash::AHashMap;
 
 use crate::build_tools::{is_strict, py_err, schema_or_config, SchemaDict};
-use crate::errors::{ErrorType, ValError, ValLineError, ValResult};
+use crate::errors::{ErrorType, LocItem, ValError, ValLineError, ValResult};
 use crate::input::{GenericMapping, Input};
 use crate::lookup_key::LookupKey;
 use crate::questions::Question;
@@ -189,10 +191,53 @@ impl Discriminator {
     }
 }
 
+#[derive(Debug, Clone, Eq, PartialEq, Hash)]
+enum ChoiceKey {
+    Int(i64),
+    Str(String),
+}
+
+impl ChoiceKey {
+    fn from_py(raw: &PyAny) -> PyResult<Self> {
+        if let Ok(py_int) = raw.extract::<i64>() {
+            Ok(Self::Int(py_int))
+        } else if let Ok(py_str) = raw.downcast::<PyString>() {
+            Ok(Self::Str(py_str.to_str()?.to_string()))
+        } else {
+            py_err!(PyTypeError; "Expected int or str, got {}", raw.get_type().name().unwrap_or("<unknown python object>"))
+        }
+    }
+
+    fn repr(&self) -> String {
+        match self {
+            Self::Int(i) => i.to_string(),
+            Self::Str(s) => format!("'{s}'"),
+        }
+    }
+}
+
+impl fmt::Display for ChoiceKey {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::Int(i) => write!(f, "{i}"),
+            Self::Str(s) => write!(f, "{s}"),
+        }
+    }
+}
+
+impl From<&ChoiceKey> for LocItem {
+    fn from(key: &ChoiceKey) -> Self {
+        match key {
+            ChoiceKey::Str(s) => s.as_str().into(),
+            ChoiceKey::Int(i) => (*i).into(),
+        }
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct TaggedUnionValidator {
-    choices: AHashMap<String, CombinedValidator>,
-    repeat_choices: Option<AHashMap<String, String>>,
+    choices: AHashMap<ChoiceKey, CombinedValidator>,
+    repeat_choices: Option<AHashMap<ChoiceKey, ChoiceKey>>,
     discriminator: Discriminator,
     from_attributes: bool,
     strict: bool,
@@ -216,25 +261,27 @@ impl BuildValidator for TaggedUnionValidator {
 
         let schema_choices: &PyDict = schema.get_as_req(intern!(py, "choices"))?;
         let mut choices = AHashMap::with_capacity(schema_choices.len());
-        let mut repeat_choices_vec: Vec<(String, String)> = Vec::new();
+        let mut repeat_choices_vec: Vec<(ChoiceKey, ChoiceKey)> = Vec::new();
         let mut first = true;
         let mut tags_repr = String::with_capacity(50);
         let mut descr = String::with_capacity(50);
 
         for (key, value) in schema_choices {
-            let tag: String = key.extract()?;
-            if let Ok(py_str) = value.downcast::<PyString>() {
-                let repeat_tag = py_str.to_str()?.to_string();
+            let tag = ChoiceKey::from_py(key)?;
+
+            if let Ok(repeat_tag) = ChoiceKey::from_py(value) {
                 repeat_choices_vec.push((tag, repeat_tag));
                 continue;
             }
+
             let validator = build_validator(value, config, build_context)?;
+            let tag_repr = tag.repr();
             if first {
                 first = false;
-                write!(tags_repr, "'{tag}'").unwrap();
+                write!(tags_repr, "{tag_repr}").unwrap();
                 descr.push_str(validator.get_name());
             } else {
-                write!(tags_repr, ", '{tag}'").unwrap();
+                write!(tags_repr, ", {tag_repr}").unwrap();
                 // no spaces in get_name() output to make loc easy to read
                 write!(descr, ",{}", validator.get_name()).unwrap();
             }
@@ -246,9 +293,10 @@ impl BuildValidator for TaggedUnionValidator {
             let mut wrong_values = Vec::with_capacity(repeat_choices_vec.len());
             let mut repeat_choices = AHashMap::with_capacity(repeat_choices_vec.len());
             for (tag, repeat_tag) in repeat_choices_vec {
-                match choices.get(repeat_tag.as_str()) {
+                match choices.get(&repeat_tag) {
                     Some(validator) => {
-                        write!(tags_repr, ", '{tag}'").unwrap();
+                        let tag_repr = tag.repr();
+                        write!(tags_repr, ", {tag_repr}").unwrap();
                         write!(descr, ",{}", validator.get_name()).unwrap();
                         repeat_choices.insert(tag, repeat_tag);
                     }
@@ -265,7 +313,7 @@ impl BuildValidator for TaggedUnionValidator {
         };
 
         let key = intern!(py, "from_attributes");
-        let from_attributes = schema_or_config(schema, config, key, key)?.unwrap_or(false);
+        let from_attributes = schema_or_config(schema, config, key, key)?.unwrap_or(true);
 
         let descr = match discriminator {
             Discriminator::SelfSchema => "self-schema".to_string(),
@@ -304,10 +352,10 @@ impl Validator for TaggedUnionValidator {
                         // errors when getting attributes which should be "raised"
                         match lookup_key.$get_method($( $dict ),+)? {
                             Some((_, value)) => {
-                                if self.strict {
-                                    value.strict_str()
+                                if let Ok(int) = value.validate_int(self.strict) {
+                                    Ok(ChoiceKey::Int(int))
                                 } else {
-                                    value.lax_str()
+                                    Ok(ChoiceKey::Str(value.validate_str(self.strict)?.as_cow()?.as_ref().to_string()))
                                 }
                             }
                             None => Err(self.tag_not_found(input)),
@@ -321,20 +369,20 @@ impl Validator for TaggedUnionValidator {
                     GenericMapping::PyMapping(mapping) => find_validator!(py_get_mapping_item, mapping),
                     GenericMapping::JsonObject(mapping) => find_validator!(json_get, mapping),
                 }?;
-                self.find_call_validator(py, tag.as_cow()?, input, extra, slots, recursion_guard)
+                self.find_call_validator(py, &tag, input, extra, slots, recursion_guard)
             }
             Discriminator::Function(ref func) => {
                 let tag = func.call1(py, (input.to_object(py),))?;
                 if tag.is_none(py) {
                     Err(self.tag_not_found(input))
                 } else {
-                    let tag: &PyString = tag.downcast(py)?;
-                    self.find_call_validator(py, tag.to_string_lossy(), input, extra, slots, recursion_guard)
+                    let tag: &PyAny = tag.downcast(py)?;
+                    self.find_call_validator(py, &(ChoiceKey::from_py(tag)?), input, extra, slots, recursion_guard)
                 }
             }
             Discriminator::SelfSchema => self.find_call_validator(
                 py,
-                self.self_schema_tag(py, input)?,
+                &ChoiceKey::Str(self.self_schema_tag(py, input)?.into_owned()),
                 input,
                 extra,
                 slots,
@@ -407,23 +455,23 @@ impl TaggedUnionValidator {
     fn find_call_validator<'s, 'data>(
         &'s self,
         py: Python<'data>,
-        tag: Cow<str>,
+        tag: &ChoiceKey,
         input: &'data impl Input<'data>,
         extra: &Extra,
         slots: &'data [CombinedValidator],
         recursion_guard: &'s mut RecursionGuard,
     ) -> ValResult<'data, PyObject> {
-        if let Some(validator) = self.choices.get(tag.as_ref()) {
+        if let Some(validator) = self.choices.get(tag) {
             return match validator.validate(py, input, extra, slots, recursion_guard) {
                 Ok(res) => Ok(res),
-                Err(err) => Err(err.with_outer_location(tag.as_ref().into())),
+                Err(err) => Err(err.with_outer_location(tag.into())),
             };
         } else if let Some(ref repeat_choices) = self.repeat_choices {
-            if let Some(choice_tag) = repeat_choices.get(tag.as_ref()) {
+            if let Some(choice_tag) = repeat_choices.get(tag) {
                 let validator = &self.choices[choice_tag];
                 return match validator.validate(py, input, extra, slots, recursion_guard) {
                     Ok(res) => Ok(res),
-                    Err(err) => Err(err.with_outer_location(tag.as_ref().into())),
+                    Err(err) => Err(err.with_outer_location(tag.into())),
                 };
             }
         }