fix: make Decimal type work with StrEnum when strict mode is not enabled

Author: vai-mikkoleppanen

src/validators/enum_.rs

@@ -169,7 +169,7 @@ impl EnumValidateValue for PlainEnumValidator {
         lookup: &LiteralLookup<PyObject>,
         strict: bool,
     ) -> ValResult<Option<PyObject>> {
-        match lookup.validate(py, input)? {
+        match lookup.validate(py, input, strict)? {
             Some((_, v)) => Ok(Some(v.clone_ref(py))),
             None => {
                 if !strict {

src/validators/literal.rs

@@ -105,6 +105,7 @@ impl<T: Debug> LiteralLookup<T> {
         &self,
         py: Python<'py>,
         input: &'a I,
+        strict: bool,
     ) -> ValResult<Option<(&'a I, &T)>> {
         if let Some(expected_bool) = &self.expected_bool {
             if let Ok(bool_value) = input.validate_bool(true) {
@@ -124,6 +125,10 @@ impl<T: Debug> LiteralLookup<T> {
                     return Ok(Some((input, &self.values[*id])));
                 }
             }
+            // if the input is a Decimal type, we need to check if its value is in the expected_ints
+            if let Ok(Some(v)) = self.try_from_dec_to_int(py, input, expected_ints) {
+                return Ok(Some(v));
+            }
         }
 
         if let Some(expected_strings) = &self.expected_str {
@@ -144,6 +149,12 @@ impl<T: Debug> LiteralLookup<T> {
                     return Ok(Some((input, &self.values[*id])));
                 }
             }
+            if !strict {
+                // if the input is a Decimal type, we need to check if its value is in the expected_strings
+                if let Ok(Some(v)) = self.try_from_dec_to_str(py, input, expected_strings) {
+                    return Ok(Some(v));
+                }
+            }
         }
         if let Some(expected_py_dict) = &self.expected_py_dict {
             // We don't use ? to unpack the result of `get_item` in the next line because unhashable
@@ -166,17 +177,14 @@ impl<T: Debug> LiteralLookup<T> {
             }
         };
 
-        // if the input is a Decimal type, we need to check if its value is in the expected_ints
-        if let Ok(Some(v)) = self.validate_decimal(py, input) {
-            return Ok(Some(v));
-        }
         Ok(None)
     }
 
-    fn validate_decimal<'a, 'py, I: Input<'py> + ?Sized>(
+    fn try_from_dec_to_int<'a, 'py, I: Input<'py> + ?Sized>(
         &self,
         py: Python<'py>,
         input: &'a I,
+        expected_ints: &AHashMap<i64, usize>,
     ) -> ValResult<Option<(&'a I, &T)>> {
         let Some(py_input) = input.as_python() else {
             return Ok(None);
@@ -186,10 +194,6 @@ impl<T: Debug> LiteralLookup<T> {
             return Ok(None);
         }
 
-        let Some(expected_ints) = &self.expected_int else {
-            return Ok(None);
-        };
-
         let Ok(EitherInt::Py(dec_value)) = decimal_as_int(input, py_input) else {
             return Ok(None);
         };
@@ -202,9 +206,39 @@ impl<T: Debug> LiteralLookup<T> {
         let Some(id) = expected_ints.get(&int) else {
             return Ok(None);
         };
+
         Ok(Some((input, &self.values[*id])))
     }
 
+    fn try_from_dec_to_str<'a, 'py, I: Input<'py> + ?Sized>(
+        &self,
+        py: Python<'py>,
+        input: &'a I,
+        expected_strings: &AHashMap<String, usize>,
+    ) -> ValResult<Option<(&'a I, &T)>> {
+        let Some(py_input) = input.as_python() else {
+            return Ok(None);
+        };
+
+        if let Ok(false) = py_input.is_instance(get_decimal_type(py)) {
+            return Ok(None);
+        }
+
+        let Ok(EitherInt::Py(dec_value)) = decimal_as_int(input, py_input) else {
+            return Ok(None);
+        };
+
+        let Ok(either_int) = dec_value.exact_int() else {
+            return Ok(None);
+        };
+        let int = either_int.into_i64(py)?;
+        if let Some(id) = expected_strings.get(&int.to_string()) {
+            return Ok(Some((input, &self.values[*id])));
+        }
+
+        Ok(None)
+    }
+
     /// Used by int enums
     pub fn validate_int<'a, 'py, I: Input<'py> + ?Sized>(
         &self,
@@ -308,7 +342,7 @@ impl Validator for LiteralValidator {
         input: &(impl Input<'py> + ?Sized),
         _state: &mut ValidationState<'_, 'py>,
     ) -> ValResult<PyObject> {
-        match self.lookup.validate(py, input)? {
+        match self.lookup.validate(py, input, _state.strict_or(false))? {
             Some((_, v)) => Ok(v.clone()),
             None => Err(ValError::new(
                 ErrorType::LiteralError {

src/validators/union.rs

@@ -464,7 +464,7 @@ impl TaggedUnionValidator {
         input: &(impl Input<'py> + ?Sized),
         state: &mut ValidationState<'_, 'py>,
     ) -> ValResult<PyObject> {
-        if let Ok(Some((tag, validator))) = self.lookup.validate(py, tag) {
+        if let Ok(Some((tag, validator))) = self.lookup.validate(py, tag, state.strict_or(false)) {
             return match validator.validate(py, input, state) {
                 Ok(res) => Ok(res),
                 Err(err) => Err(err.with_outer_location(tag)),

tests/validators/test_enums.py

@@ -1,7 +1,7 @@
 import re
 import sys
 from decimal import Decimal
-from enum import Enum, IntEnum, IntFlag
+from enum import Enum, IntEnum, IntFlag, StrEnum
 
 import pytest
 
@@ -352,31 +352,100 @@ class ColorEnum(IntEnum):
     [-1, 0, 1],
 )
 def test_enum_int_validation_should_succeed_for_decimal(value: int):
+    # GIVEN
     class MyEnum(Enum):
         VALUE = value
 
+    class MyIntEnum(IntEnum):
+        VALUE = value
+
+    # WHEN
     v = SchemaValidator(
         core_schema.with_default_schema(
             schema=core_schema.enum_schema(MyEnum, list(MyEnum.__members__.values())),
             default=MyEnum.VALUE,
         )
     )
+
+    v_int = SchemaValidator(
+        core_schema.with_default_schema(
+            schema=core_schema.enum_schema(MyIntEnum, list(MyIntEnum.__members__.values())),
+            default=MyIntEnum.VALUE,
+        )
+    )
+
+    # THEN
     assert v.validate_python(Decimal(value)) is MyEnum.VALUE
     assert v.validate_python(Decimal(float(value))) is MyEnum.VALUE
 
+    assert v_int.validate_python(Decimal(value)) is MyIntEnum.VALUE
+    assert v_int.validate_python(Decimal(float(value))) is MyIntEnum.VALUE
+
+
+def test_enum_str_validation_should_succeed_for_decimal_with_strict_disabled():
+    # GIVEN
+    class MyEnum(StrEnum):
+        VALUE = '1'
+
+    # WHEN
+    v = SchemaValidator(
+        core_schema.with_default_schema(
+            schema=core_schema.enum_schema(MyEnum, list(MyEnum.__members__.values())),
+            default=MyEnum.VALUE,
+        )
+    )
+
+    # THEN
+    assert v.validate_python(Decimal(1)) is MyEnum.VALUE
+
+
+def test_enum_str_validation_should_fail_for_decimal_with_strict_enabled():
+    # GIVEN
+    class MyEnum(StrEnum):
+        VALUE = '1'
+
+    # WHEN
+    v = SchemaValidator(
+        core_schema.with_default_schema(
+            schema=core_schema.enum_schema(MyEnum, list(MyEnum.__members__.values()), strict=True),
+            default=MyEnum.VALUE,
+        )
+    )
+
+    # THEN
+    with pytest.raises(ValidationError):
+        v.validate_python(Decimal(1))
+
 
 def test_enum_int_validation_should_fail_for_incorrect_decimal_value():
+    # GIVEN
     class MyEnum(Enum):
         VALUE = 1
 
+    class MyStrEnum(StrEnum):
+        VALUE = '2'
+
+    # WHEN
     v = SchemaValidator(
         core_schema.with_default_schema(
             schema=core_schema.enum_schema(MyEnum, list(MyEnum.__members__.values())),
             default=MyEnum.VALUE,
         )
     )
+
+    v_str = SchemaValidator(
+        core_schema.with_default_schema(
+            schema=core_schema.enum_schema(MyStrEnum, list(MyStrEnum.__members__.values())),
+            default=MyStrEnum.VALUE,
+        )
+    )
+
+    # THEN
     with pytest.raises(ValidationError):
         v.validate_python(Decimal(2))
 
     with pytest.raises(ValidationError):
         v.validate_python((1, 2))
+
+    with pytest.raises(ValidationError):
+        v_str.validate_python(Decimal(1))