Fix enum value literals (#544)

pydantic · Apr 15, 2023 · 02ec941 · 02ec941
1 parent 2e2bd34
commit 02ec941
Show file tree

Hide file tree

Showing 7 changed files with 286 additions and 326 deletions.
diff --git a/src/input/input_abstract.rs b/src/input/input_abstract.rs
@@ -90,6 +90,8 @@ pub trait Input<'a>: fmt::Debug + ToPyObject {
         self.strict_str()
     }
 
+    fn as_str_strict(&self) -> Option<&str>;
+
     fn validate_bytes(&'a self, strict: bool) -> ValResult<EitherBytes<'a>> {
         if strict {
             self.strict_bytes()
@@ -129,6 +131,8 @@ pub trait Input<'a>: fmt::Debug + ToPyObject {
         self.strict_int()
     }
 
+    fn as_int_strict(&self) -> Option<i64>;
+
     fn validate_float(&self, strict: bool, ultra_strict: bool) -> ValResult<f64> {
         if ultra_strict {
             self.ultra_strict_float()

diff --git a/src/input/input_json.rs b/src/input/input_json.rs
@@ -92,6 +92,13 @@ impl<'a> Input<'a> for JsonInput {
         }
     }
 
+    fn as_str_strict(&self) -> Option<&str> {
+        match self {
+            JsonInput::String(s) => Some(s.as_str()),
+            _ => None,
+        }
+    }
+
     fn validate_bytes(&'a self, _strict: bool) -> ValResult<EitherBytes<'a>> {
         match self {
             JsonInput::String(s) => Ok(s.as_bytes().into()),
@@ -141,6 +148,13 @@ impl<'a> Input<'a> for JsonInput {
         }
     }
 
+    fn as_int_strict(&self) -> Option<i64> {
+        match self {
+            JsonInput::Int(i) => Some(*i),
+            _ => None,
+        }
+    }
+
     fn ultra_strict_float(&self) -> ValResult<f64> {
         match self {
             JsonInput::Float(f) => Ok(*f),
@@ -349,6 +363,10 @@ impl<'a> Input<'a> for String {
         self.validate_str(false)
     }
 
+    fn as_str_strict(&self) -> Option<&str> {
+        Some(self.as_str())
+    }
+
     fn validate_bytes(&'a self, _strict: bool) -> ValResult<EitherBytes<'a>> {
         Ok(self.as_bytes().into())
     }
@@ -374,6 +392,10 @@ impl<'a> Input<'a> for String {
         }
     }
 
+    fn as_int_strict(&self) -> Option<i64> {
+        None
+    }
+
     #[cfg_attr(has_no_coverage, no_coverage)]
     fn ultra_strict_float(&self) -> ValResult<f64> {
         self.strict_float()

diff --git a/src/input/input_python.rs b/src/input/input_python.rs
@@ -220,6 +220,14 @@ impl<'a> Input<'a> for PyAny {
         }
     }
 
+    fn as_str_strict(&self) -> Option<&str> {
+        if self.get_type().is(get_py_str_type(self.py())) {
+            self.extract().ok()
+        } else {
+            None
+        }
+    }
+
     fn strict_bytes(&'a self) -> ValResult<EitherBytes<'a>> {
         if let Ok(py_bytes) = self.downcast::<PyBytes>() {
             Ok(py_bytes.into())
@@ -289,6 +297,14 @@ impl<'a> Input<'a> for PyAny {
         }
     }
 
+    fn as_int_strict(&self) -> Option<i64> {
+        if self.get_type().is(get_py_int_type(self.py())) {
+            self.extract().ok()
+        } else {
+            None
+        }
+    }
+
     fn ultra_strict_float(&self) -> ValResult<f64> {
         if matches!(self.is_instance_of::<PyInt>(), Ok(true)) {
             Err(ValError::new(ErrorType::FloatType, self))
@@ -697,3 +713,15 @@ pub fn list_as_tuple(list: &PyList) -> &PyTuple {
     };
     py_tuple.into_ref(list.py())
 }
+
+static PY_INT_TYPE: GILOnceCell<PyObject> = GILOnceCell::new();
+
+fn get_py_int_type(py: Python) -> &PyObject {
+    PY_INT_TYPE.get_or_init(py, || PyInt::type_object(py).into())
+}
+
+static PY_STR_TYPE: GILOnceCell<PyObject> = GILOnceCell::new();
+
+fn get_py_str_type(py: Python) -> &PyObject {
+    PY_STR_TYPE.get_or_init(py, || PyString::type_object(py).into())
+}