fix big int subclass extraction (pydantic#919)

alexdrydew · Aug 23, 2023 · e34a5eb · e34a5eb
1 parent e414a12
commit e34a5eb
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 4 deletions.
diff --git a/src/input/input_python.rs b/src/input/input_python.rs
@@ -304,8 +304,7 @@ impl<'a> Input<'a> for PyAny {
                 Err(ValError::new(ErrorTypeDefaults::IntType, self))
             } else {
                 // force to an int to upcast to a pure python int
-                let int = self.extract::<i64>()?;
-                Ok(EitherInt::I64(int))
+                EitherInt::upcast(self)
             }
         } else {
             Err(ValError::new(ErrorTypeDefaults::IntType, self))
@@ -320,8 +319,7 @@ impl<'a> Input<'a> for PyAny {
             str_as_int(self, &cow_str)
         } else if PyInt::is_type_of(self) {
             // force to an int to upcast to a pure python int to maintain current behaviour
-            let int = self.extract::<i64>()?;
-            Ok(EitherInt::I64(int))
+            EitherInt::upcast(self)
         } else if let Ok(float) = self.extract::<f64>() {
             float_as_int(self, float)
         } else {

diff --git a/src/input/return_enums.rs b/src/input/return_enums.rs
@@ -820,6 +820,15 @@ pub enum EitherInt<'a> {
 }
 
 impl<'a> EitherInt<'a> {
+    pub fn upcast(py_any: &'a PyAny) -> ValResult<Self> {
+        // Safety: we know that py_any is a python int
+        if let Ok(int_64) = py_any.extract::<i64>() {
+            Ok(Self::I64(int_64))
+        } else {
+            let big_int: BigInt = py_any.extract()?;
+            Ok(Self::BigInt(big_int))
+        }
+    }
     pub fn into_i64(self, py: Python<'a>) -> ValResult<'a, i64> {
         match self {
             EitherInt::I64(i) => Ok(i),

diff --git a/tests/validators/test_int.py b/tests/validators/test_int.py
@@ -421,7 +421,9 @@ def test_int_subclass() -> None:
     assert type(v_strict) == int
 
     assert v.validate_python(IntSubclass(1136885225876639845)) == 1136885225876639845
+    assert v.validate_python(IntSubclass(i64_max + 7)) == i64_max + 7
     assert v.validate_python(IntSubclass(1136885225876639845), strict=True) == 1136885225876639845
+    assert v.validate_python(IntSubclass(i64_max + 7), strict=True) == i64_max + 7
 
 
 def test_int_subclass_constraint() -> None: