refactor: make is_equal_to function generic and move it out of try_validate_any

refactor: add more test coverage

Author: vai-mikkoleppanen

src/validators/literal.rs

@@ -224,13 +224,10 @@ impl<T: Debug> LiteralLookup<T> {
             return Ok(None);
         };
 
-        let is_equal = |k: &i64| -> PyResult<bool> {
-            let equality = py_input.call_method1("__eq__", (*k,))?;
-            equality.extract::<bool>()
-        };
-
         if let Some(expected_ints) = &self.expected_int {
-            let id = expected_ints.iter().find(|(k, _)| is_equal(k).unwrap_or(false));
+            let id = expected_ints
+                .iter()
+                .find(|(&k, _)| is_equal_to(py_input, k).unwrap_or(false));
 
             if let Some((_, id)) = id {
                 return Ok(Some(&self.values[*id]));
@@ -241,10 +238,20 @@ impl<T: Debug> LiteralLookup<T> {
             return Ok(None);
         };
 
+        // try with raw strings
+        let id = expected_strings
+            .iter()
+            .find(|(k, _)| is_equal_to(py_input, k.as_str()).unwrap_or(false));
+
+        if let Some((_, id)) = id {
+            return Ok(Some(&self.values[*id]));
+        }
+
+        // try with converting to int
         let id = expected_strings
             .iter()
             .filter_map(|(k, id)| k.parse::<i64>().ok().map(|k_as_int| (k_as_int, id)))
-            .find(|(k, _)| is_equal(k).unwrap_or(false));
+            .find(|(k, _)| is_equal_to(py_input, *k).unwrap_or(false));
 
         if let Some((_, id)) = id {
             return Ok(Some(&self.values[*id]));
@@ -253,6 +260,11 @@ impl<T: Debug> LiteralLookup<T> {
     }
 }
 
+fn is_equal_to<TValue: IntoPy<Py<PyAny>>>(input: &Bound<PyAny>, value: TValue) -> PyResult<bool> {
+    let equality = input.call_method1("__eq__", (value,))?;
+    equality.extract::<bool>()
+}
+
 impl<T: PyGcTraverse + Debug> PyGcTraverse for LiteralLookup<T> {
     fn py_gc_traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
         self.expected_py_dict.py_gc_traverse(visit)?;

tests/validators/test_enums.py

@@ -384,7 +384,7 @@ class MyIntEnum(IntEnum):
 
 def test_enum_int_validation_should_succeed_for_custom_type():
     # GIVEN
-    class IntWrapper:
+    class AnyWrapper:
         def __init__(self, value):
             self.value = value
 
@@ -394,6 +394,7 @@ def __eq__(self, other: object) -> bool:
     class MyEnum(Enum):
         VALUE = 999
         SECOND_VALUE = 1000000
+        THIRD_VALUE = 'Py03'
 
     # WHEN
     v = SchemaValidator(
@@ -404,8 +405,9 @@ class MyEnum(Enum):
     )
 
     # THEN
-    assert v.validate_python(IntWrapper(999)) is MyEnum.VALUE
-    assert v.validate_python(IntWrapper(1000000)) is MyEnum.SECOND_VALUE
+    assert v.validate_python(AnyWrapper(999)) is MyEnum.VALUE
+    assert v.validate_python(AnyWrapper(1000000)) is MyEnum.SECOND_VALUE
+    assert v.validate_python(AnyWrapper('Py03')) is MyEnum.THIRD_VALUE
 
 
 def test_enum_str_validation_should_succeed_for_decimal_with_strict_disabled():
@@ -481,3 +483,25 @@ class MyStrEnum(Enum):
 
     with pytest.raises(ValidationError):
         v_str.validate_python(Decimal(2.1))
+
+
+def test_enum_int_validation_should_fail_for_plain_type_without_eq_checking():
+    # GIVEN
+    class MyEnum(Enum):
+        VALUE = 1
+
+    class MyClass:
+        def __init__(self, value):
+            self.value = value
+
+    # WHEN
+    v = SchemaValidator(
+        core_schema.with_default_schema(
+            schema=core_schema.enum_schema(MyEnum, list(MyEnum.__members__.values())),
+            default=MyEnum.VALUE,
+        )
+    )
+
+    # THEN
+    with pytest.raises(ValidationError):
+        v.validate_python(MyClass(1))

tests/validators/test_model.py

@@ -1,7 +1,7 @@
 import re
 from copy import deepcopy
 from decimal import Decimal
-from typing import Any, Callable, Dict, List, Set, Tuple
+from typing import Any, Callable, Dict, Iterable, List, Set, Tuple
 
 import pytest
 from dirty_equals import HasRepr, IsInstance
@@ -1379,3 +1379,34 @@ class MyModel:
     v.validate_assignment(m, 'enum_field_3', IntWrappable(3))
     v.validate_assignment(m, 'enum_field_4', Decimal(4))
     v.validate_assignment(m, 'enum_field_4', IntWrappable(4))
+
+
+def test_model_bug():
+    class MyModel:
+        __slots__ = (
+            '__dict__',
+            '__pydantic_fields_set__',
+            '__pydantic_extra__',
+            '__pydantic_private__',
+        )
+        x: Iterable[int]
+
+    # WHEN
+    v = SchemaValidator(
+        core_schema.model_schema(
+            MyModel,
+            core_schema.model_fields_schema(
+                {
+                    'x': core_schema.model_field(core_schema.generator_schema()),
+                },
+            ),
+        )
+    )
+    print(v)
+
+    # THEN
+    # v.validate_json('{"x": [1, 2, 3]}')
+    m = v.validate_python({'x': [1, 2, 3]})
+    print(m)
+    print(m.x)
+    print(type(m.x))