use dict structure

Author: sydney-runkle

src/validators/literal.rs

@@ -36,7 +36,8 @@ pub struct LiteralLookup<T: Debug> {
     // Catch all for unhashable types like list
     expected_py_values: Option<Vec<(Py<PyAny>, usize)>>,
     // Fallback for ints, bools, and strings to use Python hash and equality checks
-    expected_py_primitives: Option<Vec<(Py<PyAny>, usize)>>,
+    // which we can't mix with `expected_py_dict`, see tests/test_validators/test_literal.py::test_mix_int_enum_with_int
+    expected_py_primitives: Option<Py<PyDict>>,
 
     pub values: Vec<T>,
 }
@@ -48,7 +49,7 @@ impl<T: Debug> LiteralLookup<T> {
         let mut expected_str: AHashMap<String, usize> = AHashMap::new();
         let expected_py_dict = PyDict::new_bound(py);
         let mut expected_py_values = Vec::new();
-        let mut expected_py_primitives = Vec::new();
+        let expected_py_primitives = PyDict::new_bound(py);
         let mut values = Vec::new();
         for (k, v) in expected {
             let id = values.len();
@@ -60,12 +61,12 @@ impl<T: Debug> LiteralLookup<T> {
                 } else {
                     expected_bool.false_id = Some(id);
                 }
-                expected_py_primitives.push((k.as_unbound().clone_ref(py), id));
+                expected_py_primitives.set_item(&k, id)?;
             }
             if k.is_exact_instance_of::<PyInt>() {
                 if let Ok(int_64) = k.extract::<i64>() {
                     expected_int.insert(int_64, id);
-                    expected_py_primitives.push((k.as_unbound().clone_ref(py), id));
+                    expected_py_primitives.set_item(&k, id)?;
                 } else {
                     // cover the case of an int that's > i64::MAX etc.
                     expected_py_dict.set_item(k, id)?;
@@ -75,7 +76,7 @@ impl<T: Debug> LiteralLookup<T> {
                     .as_cow()
                     .map_err(|_| py_schema_error_type!("error extracting str {:?}", k))?;
                 expected_str.insert(str.to_string(), id);
-                expected_py_primitives.push((k.as_unbound().clone_ref(py), id));
+                expected_py_primitives.set_item(&k, id)?;
             } else if expected_py_dict.set_item(&k, id).is_err() {
                 expected_py_values.push((k.as_unbound().clone_ref(py), id));
             }
@@ -88,7 +89,7 @@ impl<T: Debug> LiteralLookup<T> {
             expected_str: (!expected_str.is_empty()).then_some(expected_str),
             expected_py_dict: (!expected_py_dict.is_empty()).then_some(expected_py_dict.into()),
             expected_py_values: (!expected_py_values.is_empty()).then_some(expected_py_values),
-            expected_py_primitives: (!expected_py_primitives.is_empty()).then_some(expected_py_primitives),
+            expected_py_primitives: (!expected_py_primitives.is_empty()).then_some(expected_py_primitives.into()),
             values,
         })
     }
@@ -157,20 +158,16 @@ impl<T: Debug> LiteralLookup<T> {
             }
         };
 
+        // this one must be last to avoid conflicts with the other lookups, think of this
+        // almost as a lax fallback
         if let Some(expected_py_primitives) = &self.expected_py_primitives {
             let py_input = py_input.get_or_insert_with(|| input.to_object(py));
-            let py_input_bound = py_input.bind(py);
-
-            for (k, id) in expected_py_primitives {
-                let bound_k = k.bind(py);
-                if bound_k.eq(&*py_input).unwrap_or(false) {
-                    match (bound_k.hash(), py_input_bound.hash()) {
-                        (Ok(k_hash), Ok(input_hash)) if k_hash == input_hash => {
-                            return Ok(Some((input, &self.values[*id])));
-                        }
-                        _ => continue, // Skip to the next item on hash failure or mismatch
-                    }
-                }
+            // We don't use ? to unpack the result of `get_item` in the next line because unhashable
+            // inputs will produce a TypeError, which in this case we just want to treat equivalently
+            // to a failed lookup
+            if let Ok(Some(v)) = expected_py_primitives.bind(py).get_item(&*py_input) {
+                let id: usize = v.extract().unwrap();
+                return Ok(Some((input, &self.values[id])));
             }
         };
         Ok(None)