Use PyInt for inequality and multiple of checks (#634)

Author: samuelcolvin

Cargo.lock

@@ -40,6 +40,7 @@ url = "2.3.1"
 # idna is already required by url, added here to be explicit
 idna = "0.3.0"
 base64 = "0.13.1"
+num-bigint = "0.4.3"
 
 [lib]
 name = "_pydantic_core"
@@ -50,7 +51,7 @@ crate-type = ["cdylib", "rlib"]
 extension-module = ["pyo3/extension-module"]
 # required for cargo bench
 auto-initialize = ["pyo3/auto-initialize"]
-default = ["mimalloc", "mimalloc/local_dynamic_tls", "pyo3/generate-import-lib"]
+default = ["mimalloc", "mimalloc/local_dynamic_tls", "pyo3/generate-import-lib", "pyo3/num-bigint"]
 
 [profile.release]
 lto = "fat"

Cargo.toml

@@ -3772,8 +3772,8 @@ def definition_reference_schema(
     'bool_parsing',
     'int_type',
     'int_parsing',
+    'int_parsing_size',
     'int_from_float',
-    'int_overflow',
     'float_type',
     'float_parsing',
     'bytes_type',

pydantic_core/core_schema.py

@@ -2,6 +2,7 @@ use std::borrow::Cow;
 use std::fmt;
 
 use ahash::AHashMap;
+use num_bigint::BigInt;
 use pyo3::exceptions::{PyKeyError, PyTypeError, PyValueError};
 use pyo3::once_cell::GILOnceCell;
 use pyo3::prelude::*;
@@ -174,8 +175,8 @@ pub enum ErrorType {
     // int errors
     IntType,
     IntParsing,
+    IntParsingSize,
     IntFromFloat,
-    IntOverflow,
     // ---------------------
     // float errors
     FloatType,
@@ -489,7 +490,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::IntParsingSize => "Unable to parse input string as an integer, exceed maximum size",
             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",
@@ -699,6 +700,7 @@ impl ErrorType {
 #[derive(Clone, Debug)]
 pub enum Number {
     Int(i64),
+    BigInt(BigInt),
     Float(f64),
     String(String),
 }
@@ -715,6 +717,12 @@ impl From<i64> for Number {
     }
 }
 
+impl From<BigInt> for Number {
+    fn from(i: BigInt) -> Self {
+        Self::BigInt(i)
+    }
+}
+
 impl From<f64> for Number {
     fn from(f: f64) -> Self {
         Self::Float(f)
@@ -746,6 +754,7 @@ impl fmt::Display for Number {
         match self {
             Self::Float(s) => write!(f, "{s}"),
             Self::Int(i) => write!(f, "{i}"),
+            Self::BigInt(i) => write!(f, "{i}"),
             Self::String(s) => write!(f, "{s}"),
         }
     }
@@ -754,6 +763,7 @@ impl ToPyObject for Number {
     fn to_object(&self, py: Python<'_>) -> PyObject {
         match self {
             Self::Int(i) => i.into_py(py),
+            Self::BigInt(i) => i.clone().into_py(py),
             Self::Float(f) => f.into_py(py),
             Self::String(s) => s.into_py(py),
         }

src/errors/types.rs

@@ -107,7 +107,7 @@ impl<'a> Input<'a> for JsonInput {
             JsonInput::String(s) => str_as_bool(self, s),
             JsonInput::Int(int) => int_as_bool(self, *int),
             JsonInput::Float(float) => match float_as_int(self, *float) {
-                Ok(int) => int_as_bool(self, int),
+                Ok(int) => int.as_bool().ok_or_else(|| ValError::new(ErrorType::BoolParsing, self)),
                 _ => Err(ValError::new(ErrorType::BoolType, self)),
             },
             _ => Err(ValError::new(ErrorType::BoolType, self)),
@@ -116,22 +116,23 @@ impl<'a> Input<'a> for JsonInput {
 
     fn strict_int(&'a self) -> ValResult<EitherInt<'a>> {
         match self {
-            JsonInput::Int(i) => Ok(EitherInt::Rust(*i)),
+            JsonInput::Int(i) => Ok(EitherInt::I64(*i)),
+            JsonInput::Uint(u) => Ok(EitherInt::U64(*u)),
             _ => Err(ValError::new(ErrorType::IntType, self)),
         }
     }
     fn lax_int(&'a self) -> ValResult<EitherInt<'a>> {
-        let int_result = match self {
+        match self {
             JsonInput::Bool(b) => match *b {
-                true => Ok(1),
-                false => Ok(0),
+                true => Ok(EitherInt::I64(1)),
+                false => Ok(EitherInt::I64(0)),
             },
-            JsonInput::Int(i) => Ok(*i),
+            JsonInput::Int(i) => Ok(EitherInt::I64(*i)),
+            JsonInput::Uint(u) => Ok(EitherInt::U64(*u)),
             JsonInput::Float(f) => float_as_int(self, *f),
             JsonInput::String(str) => str_as_int(self, str),
             _ => Err(ValError::new(ErrorType::IntType, self)),
-        };
-        int_result.map(EitherInt::Rust)
+        }
     }
 
     fn ultra_strict_float(&self) -> ValResult<f64> {
@@ -144,6 +145,7 @@ impl<'a> Input<'a> for JsonInput {
         match self {
             JsonInput::Float(f) => Ok(*f),
             JsonInput::Int(i) => Ok(*i as f64),
+            JsonInput::Uint(u) => Ok(*u as f64),
             _ => Err(ValError::new(ErrorType::FloatType, self)),
         }
     }
@@ -155,6 +157,7 @@ impl<'a> Input<'a> for JsonInput {
             },
             JsonInput::Float(f) => Ok(*f),
             JsonInput::Int(i) => Ok(*i as f64),
+            JsonInput::Uint(u) => Ok(*u as f64),
             JsonInput::String(str) => match str.parse::<f64>() {
                 Ok(i) => Ok(i),
                 Err(_) => Err(ValError::new(ErrorType::FloatParsing, self)),
@@ -363,7 +366,7 @@ impl<'a> Input<'a> for String {
     }
     fn lax_int(&'a self) -> ValResult<EitherInt<'a>> {
         match self.parse() {
-            Ok(i) => Ok(EitherInt::Rust(i)),
+            Ok(i) => Ok(EitherInt::I64(i)),
             Err(_) => Err(ValError::new(ErrorType::IntParsing, self)),
         }
     }

src/input/input_json.rs

@@ -260,7 +260,7 @@ impl<'a> Input<'a> for PyAny {
             int_as_bool(self, int)
         } else if let Ok(float) = self.extract::<f64>() {
             match float_as_int(self, float) {
-                Ok(int) => int_as_bool(self, int),
+                Ok(int) => int.as_bool().ok_or_else(|| ValError::new(ErrorType::BoolParsing, self)),
                 _ => Err(ValError::new(ErrorType::BoolType, self)),
             }
         } else {
@@ -287,10 +287,9 @@ impl<'a> Input<'a> for PyAny {
         if PyInt::is_exact_type_of(self) {
             Ok(EitherInt::Py(self))
         } else if let Some(cow_str) = maybe_as_string(self, ErrorType::IntParsing)? {
-            let int = str_as_int(self, &cow_str)?;
-            Ok(EitherInt::Rust(int))
+            str_as_int(self, &cow_str)
         } else if let Ok(float) = self.extract::<f64>() {
-            Ok(EitherInt::Rust(float_as_int(self, float)?))
+            float_as_int(self, float)
         } else {
             Err(ValError::new(ErrorType::IntType, self))
         }

src/input/input_python.rs

@@ -12,6 +12,7 @@ pub enum JsonInput {
     Null,
     Bool(bool),
     Int(i64),
+    Uint(u64),
     Float(f64),
     String(String),
     Array(JsonArray),
@@ -26,6 +27,7 @@ impl ToPyObject for JsonInput {
             Self::Null => py.None(),
             Self::Bool(b) => b.into_py(py),
             Self::Int(i) => i.into_py(py),
+            Self::Uint(i) => i.into_py(py),
             Self::Float(f) => f.into_py(py),
             Self::String(s) => s.into_py(py),
             Self::Array(v) => PyList::new(py, v.iter().map(|v| v.to_object(py))).into_py(py),
@@ -64,7 +66,10 @@ impl<'de> Deserialize<'de> for JsonInput {
             }
 
             fn visit_u64<E>(self, value: u64) -> Result<JsonInput, E> {
-                Ok(JsonInput::Int(value as i64))
+                match i64::try_from(value) {
+                    Ok(i) => Ok(JsonInput::Int(i)),
+                    Err(_) => Ok(JsonInput::Uint(value)),
+                }
             }
 
             fn visit_f64<E>(self, value: f64) -> Result<JsonInput, E> {

src/input/parse_json.rs

@@ -1,6 +1,8 @@
 use std::borrow::Cow;
 use std::slice::Iter as SliceIter;
 
+use num_bigint::BigInt;
+
 use pyo3::prelude::*;
 use pyo3::types::iter::PyDictIterator;
 use pyo3::types::{
@@ -820,25 +822,72 @@ impl<'a> IntoPy<PyObject> for EitherBytes<'a> {
 
 #[cfg_attr(debug_assertions, derive(Debug))]
 pub enum EitherInt<'a> {
-    Rust(i64),
+    I64(i64),
+    U64(u64),
+    BigInt(BigInt),
     Py(&'a PyAny),
 }
 
-impl<'a> TryInto<i64> for EitherInt<'a> {
-    type Error = ValError<'a>;
+impl<'a> EitherInt<'a> {
+    pub fn into_i64(self, py: Python<'a>) -> ValResult<'a, i64> {
+        match self {
+            EitherInt::I64(i) => Ok(i),
+            EitherInt::U64(u) => match i64::try_from(u) {
+                Ok(u) => Ok(u),
+                Err(_) => Err(ValError::new(ErrorType::IntParsingSize, u.into_py(py).into_ref(py))),
+            },
+            EitherInt::BigInt(u) => match i64::try_from(u) {
+                Ok(u) => Ok(u),
+                Err(e) => Err(ValError::new(
+                    ErrorType::IntParsingSize,
+                    e.into_original().into_py(py).into_ref(py),
+                )),
+            },
+            EitherInt::Py(i) => i.extract().map_err(|_| ValError::new(ErrorType::IntParsingSize, i)),
+        }
+    }
+
+    pub fn as_bool(&self) -> Option<bool> {
+        match self {
+            EitherInt::I64(i) => match i {
+                0 => Some(false),
+                1 => Some(true),
+                _ => None,
+            },
+            EitherInt::U64(u) => match u {
+                0 => Some(false),
+                1 => Some(true),
+                _ => None,
+            },
+            EitherInt::BigInt(i) => match u8::try_from(i) {
+                Ok(0) => Some(false),
+                Ok(1) => Some(true),
+                _ => None,
+            },
+            EitherInt::Py(i) => match i.extract::<u8>() {
+                Ok(0) => Some(false),
+                Ok(1) => Some(true),
+                _ => None,
+            },
+        }
+    }
 
-    fn try_into(self) -> ValResult<'a, i64> {
+    pub fn as_bigint(&self) -> PyResult<BigInt> {
         match self {
-            EitherInt::Rust(i) => Ok(i),
-            EitherInt::Py(i) => i.extract().map_err(|_| ValError::new(ErrorType::IntOverflow, i)),
+            EitherInt::I64(i) => Ok(BigInt::from(*i)),
+            EitherInt::U64(u) => Ok(BigInt::from(*u)),
+            EitherInt::BigInt(i) => Ok(i.clone()),
+            EitherInt::Py(i) => i.extract(),
         }
     }
 }
 
 impl<'a> IntoPy<PyObject> for EitherInt<'a> {
     fn into_py(self, py: Python<'_>) -> PyObject {
         match self {
-            Self::Rust(int) => int.into_py(py),
+            Self::I64(int) => int.into_py(py),
+            Self::U64(int) => int.into_py(py),
+            Self::BigInt(int) => int.into_py(py),
             Self::Py(int) => int.into_py(py),
         }
     }