Int validation performance improvement (#620)

Author: samuelcolvin

pydantic_core/core_schema.py

@@ -3761,6 +3761,7 @@ def definition_reference_schema(
     'int_type',
     'int_parsing',
     'int_from_float',
+    'int_overflow',
     'float_type',
     'float_parsing',
     'bytes_type',

src/errors/types.rs

@@ -175,6 +175,7 @@ pub enum ErrorType {
     IntType,
     IntParsing,
     IntFromFloat,
+    IntOverflow,
     // ---------------------
     // float errors
     FloatType,
@@ -488,6 +489,7 @@ impl ErrorType {
             Self::IntType => "Input should be a valid integer",
             Self::IntParsing => "Input should be a valid integer, unable to parse string as an integer",
             Self::IntFromFloat => "Input should be a valid integer, got a number with a fractional part",
+            Self::IntOverflow => "Input integer too large to convert to 64-bit integer",
             Self::FloatType => "Input should be a valid number",
             Self::FloatParsing => "Input should be a valid number, unable to parse string as an number",
             Self::BytesType => "Input should be a valid bytes",

src/input/input_abstract.rs

@@ -7,7 +7,7 @@ use crate::errors::{InputValue, LocItem, ValResult};
 use crate::{PyMultiHostUrl, PyUrl};
 
 use super::datetime::{EitherDate, EitherDateTime, EitherTime, EitherTimedelta};
-use super::return_enums::{EitherBytes, EitherString};
+use super::return_enums::{EitherBytes, EitherInt, EitherString};
 use super::{GenericArguments, GenericIterable, GenericIterator, GenericMapping, JsonInput};
 
 #[derive(Debug, Clone, Copy)]
@@ -90,8 +90,6 @@ 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()
@@ -118,21 +116,19 @@ pub trait Input<'a>: fmt::Debug + ToPyObject {
         self.strict_bool()
     }
 
-    fn validate_int(&self, strict: bool) -> ValResult<i64> {
+    fn validate_int(&'a self, strict: bool) -> ValResult<EitherInt<'a>> {
         if strict {
             self.strict_int()
         } else {
             self.lax_int()
         }
     }
-    fn strict_int(&self) -> ValResult<i64>;
+    fn strict_int(&'a self) -> ValResult<EitherInt<'a>>;
     #[cfg_attr(has_no_coverage, no_coverage)]
-    fn lax_int(&self) -> ValResult<i64> {
+    fn lax_int(&'a self) -> ValResult<EitherInt<'a>> {
         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()

src/input/input_json.rs

@@ -10,8 +10,8 @@ use super::datetime::{
 use super::parse_json::JsonArray;
 use super::shared::{float_as_int, int_as_bool, map_json_err, str_as_bool, str_as_int};
 use super::{
-    EitherBytes, EitherString, EitherTimedelta, GenericArguments, GenericIterable, GenericIterator, GenericMapping,
-    Input, JsonArgs, JsonInput,
+    EitherBytes, EitherInt, EitherString, EitherTimedelta, GenericArguments, GenericIterable, GenericIterator,
+    GenericMapping, Input, JsonArgs, JsonInput,
 };
 
 impl<'a> Input<'a> for JsonInput {
@@ -84,13 +84,6 @@ 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()),
@@ -121,14 +114,14 @@ impl<'a> Input<'a> for JsonInput {
         }
     }
 
-    fn strict_int(&self) -> ValResult<i64> {
+    fn strict_int(&'a self) -> ValResult<EitherInt<'a>> {
         match self {
-            JsonInput::Int(i) => Ok(*i),
+            JsonInput::Int(i) => Ok(EitherInt::Rust(*i)),
             _ => Err(ValError::new(ErrorType::IntType, self)),
         }
     }
-    fn lax_int(&self) -> ValResult<i64> {
-        match self {
+    fn lax_int(&'a self) -> ValResult<EitherInt<'a>> {
+        let int_result = match self {
             JsonInput::Bool(b) => match *b {
                 true => Ok(1),
                 false => Ok(0),
@@ -137,14 +130,8 @@ impl<'a> Input<'a> for JsonInput {
             JsonInput::Float(f) => float_as_int(self, *f),
             JsonInput::String(str) => str_as_int(self, str),
             _ => Err(ValError::new(ErrorType::IntType, self)),
-        }
-    }
-
-    fn as_int_strict(&self) -> Option<i64> {
-        match self {
-            JsonInput::Int(i) => Some(*i),
-            _ => None,
-        }
+        };
+        int_result.map(EitherInt::Rust)
     }
 
     fn ultra_strict_float(&self) -> ValResult<f64> {
@@ -356,10 +343,6 @@ 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())
     }
@@ -375,20 +358,16 @@ impl<'a> Input<'a> for String {
         str_as_bool(self, self)
     }
 
-    fn strict_int(&self) -> ValResult<i64> {
+    fn strict_int(&'a self) -> ValResult<EitherInt<'a>> {
         Err(ValError::new(ErrorType::IntType, self))
     }
-    fn lax_int(&self) -> ValResult<i64> {
+    fn lax_int(&'a self) -> ValResult<EitherInt<'a>> {
         match self.parse() {
-            Ok(i) => Ok(i),
+            Ok(i) => Ok(EitherInt::Rust(i)),
             Err(_) => Err(ValError::new(ErrorType::IntParsing, self)),
         }
     }
 
-    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()

src/input/input_python.rs

@@ -1,7 +1,6 @@
 use std::borrow::Cow;
 use std::str::from_utf8;
 
-use pyo3::once_cell::GILOnceCell;
 use pyo3::prelude::*;
 use pyo3::types::{
     PyBool, PyByteArray, PyBytes, PyDate, PyDateTime, PyDelta, PyDict, PyFrozenSet, PyInt, PyIterator, PyList,
@@ -22,8 +21,8 @@ use super::datetime::{
 };
 use super::shared::{float_as_int, int_as_bool, map_json_err, str_as_bool, str_as_int};
 use super::{
-    py_string_str, EitherBytes, EitherString, EitherTimedelta, GenericArguments, GenericIterable, GenericIterator,
-    GenericMapping, Input, JsonInput, PyArgs,
+    py_string_str, EitherBytes, EitherInt, EitherString, EitherTimedelta, GenericArguments, GenericIterable,
+    GenericIterator, GenericMapping, Input, JsonInput, PyArgs,
 };
 
 #[cfg(not(PyPy))]
@@ -227,14 +226,6 @@ 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())
@@ -281,37 +272,34 @@ impl<'a> Input<'a> for PyAny {
         }
     }
 
-    fn strict_int(&self) -> ValResult<i64> {
-        // bool check has to come before int check as bools would be cast to ints below
-        if self.extract::<bool>().is_ok() {
-            Err(ValError::new(ErrorType::IntType, self))
-        } else if let Ok(int) = self.extract::<i64>() {
-            Ok(int)
+    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))
         }
     }
 
-    fn lax_int(&self) -> ValResult<i64> {
-        if let Ok(int) = self.extract::<i64>() {
-            Ok(int)
+    fn lax_int(&'a self) -> ValResult<EitherInt<'a>> {
+        if PyInt::is_exact_type_of(self) {
+            Ok(EitherInt::Py(self))
         } else if let Some(cow_str) = maybe_as_string(self, ErrorType::IntParsing)? {
-            str_as_int(self, &cow_str)
+            let int = str_as_int(self, &cow_str)?;
+            Ok(EitherInt::Rust(int))
         } else if let Ok(float) = self.extract::<f64>() {
-            float_as_int(self, float)
+            Ok(EitherInt::Rust(float_as_int(self, float)?))
         } else {
             Err(ValError::new(ErrorType::IntType, self))
         }
     }
 
-    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))
@@ -322,10 +310,13 @@ impl<'a> Input<'a> for PyAny {
         }
     }
     fn strict_float(&self) -> ValResult<f64> {
-        if self.extract::<bool>().is_ok() {
-            Err(ValError::new(ErrorType::FloatType, self))
-        } else if let Ok(float) = self.extract::<f64>() {
-            Ok(float)
+        if let Ok(float) = self.extract::<f64>() {
+            // bools are cast to floats as either 0.0 or 1.0, so check for bool type in this specific case
+            if (float == 0.0 || float == 1.0) && PyBool::is_exact_type_of(self) {
+                Err(ValError::new(ErrorType::FloatType, self))
+            } else {
+                Ok(float)
+            }
         } else {
             Err(ValError::new(ErrorType::FloatType, self))
         }
@@ -515,7 +506,7 @@ impl<'a> Input<'a> for PyAny {
     }
 
     fn strict_date(&self) -> ValResult<EitherDate> {
-        if self.downcast::<PyDateTime>().is_ok() {
+        if PyDateTime::is_type_of(self) {
             // have to check if it's a datetime first, otherwise the line below converts to a date
             Err(ValError::new(ErrorType::DateType, self))
         } else if let Ok(date) = self.downcast::<PyDate>() {
@@ -526,7 +517,7 @@ impl<'a> Input<'a> for PyAny {
     }
 
     fn lax_date(&self) -> ValResult<EitherDate> {
-        if self.downcast::<PyDateTime>().is_ok() {
+        if PyDateTime::is_type_of(self) {
             // have to check if it's a datetime first, otherwise the line below converts to a date
             // even if we later try coercion from a datetime, we don't want to return a datetime now
             Err(ValError::new(ErrorType::DateType, self))
@@ -558,7 +549,7 @@ impl<'a> Input<'a> for PyAny {
             bytes_as_time(self, str.as_bytes())
         } else if let Ok(py_bytes) = self.downcast::<PyBytes>() {
             bytes_as_time(self, py_bytes.as_bytes())
-        } else if self.downcast::<PyBool>().is_ok() {
+        } else if PyBool::is_exact_type_of(self) {
             Err(ValError::new(ErrorType::TimeType, self))
         } else if let Ok(int) = self.extract::<i64>() {
             int_as_time(self, int, 0)
@@ -585,7 +576,7 @@ impl<'a> Input<'a> for PyAny {
             bytes_as_datetime(self, str.as_bytes())
         } else if let Ok(py_bytes) = self.downcast::<PyBytes>() {
             bytes_as_datetime(self, py_bytes.as_bytes())
-        } else if self.downcast::<PyBool>().is_ok() {
+        } else if PyBool::is_exact_type_of(self) {
             Err(ValError::new(ErrorType::DatetimeType, self))
         } else if let Ok(int) = self.extract::<i64>() {
             int_as_datetime(self, int, 0)
@@ -661,7 +652,7 @@ fn is_builtin_str(py_str: &PyString) -> bool {
 }
 
 #[cfg(PyPy)]
-static DICT_KEYS_TYPE: GILOnceCell<Py<PyType>> = GILOnceCell::new();
+static DICT_KEYS_TYPE: pyo3::once_cell::GILOnceCell<Py<PyType>> = pyo3::once_cell::GILOnceCell::new();
 
 #[cfg(PyPy)]
 fn is_dict_keys_type(v: &PyAny) -> bool {
@@ -679,7 +670,7 @@ fn is_dict_keys_type(v: &PyAny) -> bool {
 }
 
 #[cfg(PyPy)]
-static DICT_VALUES_TYPE: GILOnceCell<Py<PyType>> = GILOnceCell::new();
+static DICT_VALUES_TYPE: pyo3::once_cell::GILOnceCell<Py<PyType>> = pyo3::once_cell::GILOnceCell::new();
 
 #[cfg(PyPy)]
 fn is_dict_values_type(v: &PyAny) -> bool {
@@ -697,7 +688,7 @@ fn is_dict_values_type(v: &PyAny) -> bool {
 }
 
 #[cfg(PyPy)]
-static DICT_ITEMS_TYPE: GILOnceCell<Py<PyType>> = GILOnceCell::new();
+static DICT_ITEMS_TYPE: pyo3::once_cell::GILOnceCell<Py<PyType>> = pyo3::once_cell::GILOnceCell::new();
 
 #[cfg(PyPy)]
 fn is_dict_items_type(v: &PyAny) -> bool {
@@ -722,15 +713,3 @@ 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())
-}

src/input/mod.rs

@@ -17,9 +17,9 @@ pub(crate) use datetime::{
 pub(crate) use input_abstract::{Input, InputType};
 pub(crate) use parse_json::{JsonInput, JsonObject};
 pub(crate) use return_enums::{
-    py_string_str, AttributesGenericIterator, DictGenericIterator, EitherBytes, EitherString, GenericArguments,
-    GenericIterable, GenericIterator, GenericMapping, JsonArgs, JsonObjectGenericIterator, MappingGenericIterator,
-    PyArgs,
+    py_string_str, AttributesGenericIterator, DictGenericIterator, EitherBytes, EitherInt, EitherString,
+    GenericArguments, GenericIterable, GenericIterator, GenericMapping, JsonArgs, JsonObjectGenericIterator,
+    MappingGenericIterator, PyArgs,
 };
 
 // Defined here as it's not exported by pyo3

src/input/return_enums.rs

@@ -817,3 +817,29 @@ impl<'a> IntoPy<PyObject> for EitherBytes<'a> {
         }
     }
 }
+
+#[cfg_attr(debug_assertions, derive(Debug))]
+pub enum EitherInt<'a> {
+    Rust(i64),
+    Py(&'a PyAny),
+}
+
+impl<'a> TryInto<i64> for EitherInt<'a> {
+    type Error = ValError<'a>;
+
+    fn try_into(self) -> ValResult<'a, i64> {
+        match self {
+            EitherInt::Rust(i) => Ok(i),
+            EitherInt::Py(i) => i.extract().map_err(|_| ValError::new(ErrorType::IntOverflow, i)),
+        }
+    }
+}
+
+impl<'a> IntoPy<PyObject> for EitherInt<'a> {
+    fn into_py(self, py: Python<'_>) -> PyObject {
+        match self {
+            Self::Rust(int) => int.into_py(py),
+            Self::Py(int) => int.into_py(py),
+        }
+    }
+}