Re-use literal validator logic for tagged unions

adriangb · adriangb · commit 89e89f4a3d04 · 2023-06-01T11:38:40.000-05:00
diff --git a/src/input/input_python.rs b/src/input/input_python.rs
@@ -271,13 +271,6 @@ impl<'a> Input<'a> for PyAny {
     fn strict_int(&'a self) -> ValResult<EitherInt<'a>> {
         if PyInt::is_exact_type_of(self) {
             Ok(EitherInt::Py(self))
-        } else if PyInt::is_type_of(self) {
-            // bools are a subclass of int, so check for bool type in this specific case
-            if PyBool::is_exact_type_of(self) {
-                Err(ValError::new(ErrorType::IntType, self))
-            } else {
-                Ok(EitherInt::Py(self))
-            }
         } else {
             Err(ValError::new(ErrorType::IntType, self))
         }
diff --git a/src/validators/literal.rs b/src/validators/literal.rs
@@ -1,7 +1,8 @@
 // Validator for things inside of a typing.Literal[]
 // which can be an int, a string, bytes or an Enum value (including `class Foo(str, Enum)` type enums)
+use core::fmt::Debug;
 
-use ahash::AHashSet;
+use ahash::AHashMap;
 use pyo3::intern;
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyList};
@@ -15,15 +16,96 @@ use crate::tools::SchemaDict;
 use super::{BuildValidator, CombinedValidator, Definitions, DefinitionsBuilder, Extra, Validator};
 
 #[derive(Debug, Clone)]
-pub struct LiteralValidator {
+pub struct LiteralLookup<T: Clone + Debug> {
     // Specialized lookups for ints and strings because they
     // (1) are easy to convert between Rust and Python
     // (2) hashing them in Rust is very fast
     // (3) are the most commonly used things in Literal[...]
-    expected_int: Option<AHashSet<i64>>,
-    expected_str: Option<AHashSet<String>>,
+    expected_int: Option<AHashMap<i64, usize>>,
+    expected_str: Option<AHashMap<String, usize>>,
     // Catch all for Enum and bytes (the latter only because it is seldom used)
     expected_py: Option<Py<PyDict>>,
+    pub values: Vec<T>,
+}
+
+impl<T: Clone + Debug> LiteralLookup<T> {
+    pub fn new<'py>(py: Python<'py>, expected: impl Iterator<Item = (&'py PyAny, T)>) -> PyResult<Self> {
+        let mut expected_int = AHashMap::new();
+        let mut expected_str = AHashMap::new();
+        let expected_py = PyDict::new(py);
+        let mut values = Vec::new();
+        for (k, v) in expected {
+            let id = values.len();
+            values.push(v);
+            if let Ok(either_int) = k.strict_int() {
+                let int = either_int
+                    .into_i64(py)
+                    .map_err(|_| py_schema_error_type!("error extracting int {:?}", k))?;
+                expected_int.insert(int, id);
+            } else if let Ok(either_str) = k.strict_str() {
+                let str = either_str
+                    .as_cow()
+                    .map_err(|_| py_schema_error_type!("error extracting str {:?}", k))?;
+                expected_str.insert(str.to_string(), id);
+            } else {
+                expected_py.set_item(k, id)?;
+            }
+        }
+
+        Ok(Self {
+            expected_int: match expected_int.is_empty() {
+                true => None,
+                false => Some(expected_int),
+            },
+            expected_str: match expected_str.is_empty() {
+                true => None,
+                false => Some(expected_str),
+            },
+            expected_py: match expected_py.is_empty() {
+                true => None,
+                false => Some(expected_py.into()),
+            },
+            values,
+        })
+    }
+
+    pub fn validate<'data, I: Input<'data>>(
+        &self,
+        py: Python<'data>,
+        input: &'data I,
+    ) -> ValResult<'data, Option<(&'data I, &T)>> {
+        // dbg!(input.to_object(py).as_ref(py).repr().unwrap());
+        if let Some(expected_ints) = &self.expected_int {
+            if let Ok(either_int) = input.strict_int() {
+                let int = either_int.into_i64(py)?;
+                if let Some(id) = expected_ints.get(&int) {
+                    return Ok(Some((input, &self.values[*id])));
+                }
+            }
+        }
+        if let Some(expected_strings) = &self.expected_str {
+            // dbg!(expected_strings);
+            if let Ok(either_str) = input.strict_str() {
+                let cow = either_str.as_cow()?;
+                if let Some(id) = expected_strings.get(cow.as_ref()) {
+                    return Ok(Some((input, &self.values[*id])));
+                }
+            }
+        }
+        // must be an enum or bytes
+        if let Some(expected_py) = &self.expected_py {
+            if let Some(v) = expected_py.as_ref(py).get_item(input) {
+                let id: usize = v.extract().unwrap();
+                return Ok(Some((input, &self.values[id])));
+            }
+        };
+        Ok(None)
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct LiteralValidator {
+    lookup: LiteralLookup<PyObject>,
     expected_repr: String,
     name: String,
 }
@@ -41,32 +123,14 @@ impl BuildValidator for LiteralValidator {
             return py_schema_err!("`expected` should have length > 0");
         }
         let py = expected.py();
-        // Literal[...] only supports int, str, bytes or enums, all of which can be hashed
-        let mut expected_int = AHashSet::new();
-        let mut expected_str = AHashSet::new();
-        let expected_py = PyDict::new(py);
         let mut repr_args: Vec<String> = Vec::new();
         for item in expected.iter() {
             repr_args.push(item.repr()?.extract()?);
-            if let Ok(either_int) = item.strict_int() {
-                let int = either_int
-                    .into_i64(py)
-                    .map_err(|_| py_schema_error_type!("error extracting int {:?}", item))?;
-                expected_int.insert(int);
-            } else if let Ok(either_str) = item.strict_str() {
-                let str = either_str
-                    .as_cow()
-                    .map_err(|_| py_schema_error_type!("error extracting str {:?}", item))?;
-                expected_str.insert(str.to_string());
-            } else {
-                expected_py.set_item(item, item)?;
-            }
         }
         let (expected_repr, name) = expected_repr_name(repr_args, "literal");
+        let lookup = LiteralLookup::new(py, expected.iter().map(|v| (v, v.to_object(py))))?;
         Ok(CombinedValidator::Literal(Self {
-            expected_int: (!expected_int.is_empty()).then_some(expected_int),
-            expected_str: (!expected_str.is_empty()).then_some(expected_str),
-            expected_py: (!expected_py.is_empty()).then_some(expected_py.into()),
+            lookup,
             expected_repr,
             name,
         }))
@@ -82,34 +146,15 @@ impl Validator for LiteralValidator {
         _definitions: &'data Definitions<CombinedValidator>,
         _recursion_guard: &'s mut RecursionGuard,
     ) -> ValResult<'data, PyObject> {
-        if let Some(expected_ints) = &self.expected_int {
-            if let Ok(either_int) = input.strict_int() {
-                let int = either_int.into_i64(py)?;
-                if expected_ints.contains(&int) {
-                    return Ok(input.to_object(py));
-                }
-            }
-        }
-        if let Some(expected_strings) = &self.expected_str {
-            if let Ok(either_str) = input.strict_str() {
-                let cow = either_str.as_cow()?;
-                if expected_strings.contains(cow.as_ref()) {
-                    return Ok(input.to_object(py));
-                }
-            }
+        match self.lookup.validate(py, input)? {
+            Some((_, v)) => Ok(v.clone()),
+            None => Err(ValError::new(
+                ErrorType::LiteralError {
+                    expected: self.expected_repr.clone(),
+                },
+                input,
+            )),
         }
-        // must be an enum or bytes
-        if let Some(expected_py) = &self.expected_py {
-            if let Some(v) = expected_py.as_ref(py).get_item(input) {
-                return Ok(v.into());
-            }
-        };
-        Err(ValError::new(
-            ErrorType::LiteralError {
-                expected: self.expected_repr.clone(),
-            },
-            input,
-        ))
     }
 
     fn different_strict_behavior(
diff --git a/src/validators/union.rs b/src/validators/union.rs
@@ -13,7 +13,7 @@ use crate::recursion_guard::RecursionGuard;
 use crate::tools::SchemaDict;
 
 use super::custom_error::CustomError;
-use super::literal::LiteralValidator;
+use super::literal::LiteralLookup;
 use super::{build_validator, BuildValidator, CombinedValidator, Definitions, DefinitionsBuilder, Extra, Validator};
 
 #[derive(Debug, Clone)]
@@ -218,10 +218,8 @@ impl Discriminator {
 
 #[derive(Debug, Clone)]
 pub struct TaggedUnionValidator {
-    choices: Py<PyDict>,
-    validators: Vec<CombinedValidator>,
     discriminator: Discriminator,
-    discriminator_validator: Box<CombinedValidator>,
+    lookup: LiteralLookup<CombinedValidator>,
     from_attributes: bool,
     strict: bool,
     custom_error: Option<CustomError>,
@@ -243,21 +241,15 @@ impl BuildValidator for TaggedUnionValidator {
         let discriminator_repr = discriminator.to_string_py(py)?;
 
         let choices = PyDict::new(py);
-        let mut validators = Vec::with_capacity(choices.len());
         let mut tags_repr = String::with_capacity(50);
         let mut descr = String::with_capacity(50);
         let mut first = true;
         let mut discriminators = Vec::with_capacity(choices.len());
         let schema_choices: &PyDict = schema.get_as_req(intern!(py, "choices"))?;
-        let schema_choice_keys = schema_choices.keys();
-        let discriminator_validator_schema = PyDict::new(py);
-        discriminator_validator_schema.set_item(intern!(py, "type"), intern!(py, "literal").as_ref())?;
-        discriminator_validator_schema.set_item(intern!(py, "expected"), schema_choice_keys.as_ref())?;
-        let discriminator_validator = LiteralValidator::build(discriminator_validator_schema, config, definitions)?;
+        let mut lookup_map = Vec::with_capacity(choices.len());
         for (choice_key, choice_schema) in schema_choices.iter() {
             discriminators.push(choice_key);
             let validator = build_validator(choice_schema, config, definitions)?;
-            choices.set_item(choice_key, validators.len())?;
             let tag_repr = choice_key.repr()?.to_string();
             if first {
                 first = false;
@@ -268,9 +260,11 @@ impl BuildValidator for TaggedUnionValidator {
                 // no spaces in get_name() output to make loc easy to read
                 write!(descr, ",{}", validator.get_name()).unwrap();
             }
-            validators.push(validator);
+            lookup_map.push((choice_key, validator));
         }
 
+        let lookup = LiteralLookup::new(py, lookup_map.into_iter())?;
+
         let key = intern!(py, "from_attributes");
         let from_attributes = schema_or_config(schema, config, key, key)?.unwrap_or(true);
 
@@ -280,10 +274,8 @@ impl BuildValidator for TaggedUnionValidator {
         };
 
         Ok(Self {
-            choices: choices.into(),
-            validators,
             discriminator,
-            discriminator_validator: Box::new(discriminator_validator),
+            lookup,
             from_attributes,
             strict: is_strict(schema, config)?,
             custom_error: CustomError::build(schema, config, definitions)?,
@@ -352,7 +344,8 @@ impl Validator for TaggedUnionValidator {
         definitions: Option<&DefinitionsBuilder<CombinedValidator>>,
         ultra_strict: bool,
     ) -> bool {
-        self.validators
+        self.lookup
+            .values
             .iter()
             .any(|v| v.different_strict_behavior(definitions, ultra_strict))
     }
@@ -362,7 +355,8 @@ impl Validator for TaggedUnionValidator {
     }
 
     fn complete(&mut self, definitions: &DefinitionsBuilder<CombinedValidator>) -> PyResult<()> {
-        self.validators
+        self.lookup
+            .values
             .iter_mut()
             .try_for_each(|validator| validator.complete(definitions))
     }
@@ -423,19 +417,11 @@ impl TaggedUnionValidator {
         definitions: &'data Definitions<CombinedValidator>,
         recursion_guard: &'s mut RecursionGuard,
     ) -> ValResult<'data, PyObject> {
-        if let Ok(tag) = self
-            .discriminator_validator
-            .validate(py, tag, extra, definitions, recursion_guard)
-        {
-            let tag = tag.as_ref(py);
-            if let Some(validator_idx) = self.choices.as_ref(py).get_item(tag) {
-                // We know this will always be a usize because we put it there ourselves
-                let validator = &self.validators[usize::extract(validator_idx).unwrap()];
-                return match validator.validate(py, input, extra, definitions, recursion_guard) {
-                    Ok(res) => Ok(res),
-                    Err(err) => Err(err.with_outer_location(LocItem::try_from(tag)?)),
-                };
-            }
+        if let Ok(Some((tag, validator))) = self.lookup.validate(py, tag) {
+            return match validator.validate(py, input, extra, definitions, recursion_guard) {
+                Ok(res) => Ok(res),
+                Err(err) => Err(err.with_outer_location(LocItem::try_from(tag.to_object(py).into_ref(py))?)),
+            };
         }
         match self.custom_error {
             Some(ref custom_error) => Err(custom_error.as_val_error(input)),
