Refactor BigInt boilerplate (#1421)

Author: sydney-runkle

src/input/return_enums.rs

@@ -20,7 +20,7 @@ use crate::errors::{
     py_err_string, ErrorType, ErrorTypeDefaults, InputValue, ToErrorValue, ValError, ValLineError, ValResult,
 };
 use crate::py_gc::PyGcTraverse;
-use crate::tools::{extract_i64, new_py_string, py_err};
+use crate::tools::{extract_i64, extract_int, new_py_string, py_err};
 use crate::validators::{CombinedValidator, Exactness, ValidationState, Validator};
 
 use super::{py_error_on_minusone, BorrowInput, Input};
@@ -662,6 +662,15 @@ pub enum Int {
     Big(BigInt),
 }
 
+impl IntoPy<PyObject> for Int {
+    fn into_py(self, py: Python<'_>) -> PyObject {
+        match self {
+            Self::I64(i) => i.into_py(py),
+            Self::Big(big_i) => big_i.into_py(py),
+        }
+    }
+}
+
 // The default serialization for BigInt is some internal representation which roundtrips efficiently
 // but is not the JSON value which users would expect to see.
 fn serialize_bigint_as_number<S>(big_int: &BigInt, serializer: S) -> Result<S::Ok, S::Error>
@@ -706,12 +715,9 @@ impl<'a> Rem for &'a Int {
 
 impl FromPyObject<'_> for Int {
     fn extract_bound(obj: &Bound<'_, PyAny>) -> PyResult<Self> {
-        if let Some(i) = extract_i64(obj) {
-            Ok(Int::I64(i))
-        } else if let Ok(b) = obj.extract::<BigInt>() {
-            Ok(Int::Big(b))
-        } else {
-            py_err!(PyTypeError; "Expected int, got {}", obj.get_type())
+        match extract_int(obj) {
+            Some(i) => Ok(i),
+            None => py_err!(PyTypeError; "Expected int, got {}", obj.get_type()),
         }
     }
 }

src/serializers/infer.rs

@@ -1,6 +1,5 @@
 use std::borrow::Cow;
 
-use num_bigint::BigInt;
 use pyo3::exceptions::PyTypeError;
 use pyo3::intern;
 use pyo3::prelude::*;
@@ -12,7 +11,7 @@ use serde::ser::{Error, Serialize, SerializeMap, SerializeSeq, Serializer};
 
 use crate::input::{EitherTimedelta, Int};
 use crate::serializers::type_serializers;
-use crate::tools::{extract_i64, py_err, safe_repr};
+use crate::tools::{extract_int, py_err, safe_repr};
 use crate::url::{PyMultiHostUrl, PyUrl};
 
 use super::config::InfNanMode;
@@ -118,10 +117,8 @@ pub(crate) fn infer_to_python_known(
             ObType::None | ObType::Bool | ObType::Int | ObType::Str => value.into_py(py),
             // have to do this to make sure subclasses of for example str are upcast to `str`
             ObType::IntSubclass => {
-                if let Some(v) = extract_i64(value) {
-                    v.into_py(py)
-                } else if let Ok(b) = value.extract::<BigInt>() {
-                    b.into_py(py)
+                if let Some(i) = extract_int(value) {
+                    i.into_py(py)
                 } else {
                     return py_err!(PyTypeError; "Expected int, got {}", safe_repr(value));
                 }

src/tools.rs

@@ -1,10 +1,13 @@
 use core::fmt;
 
+use num_bigint::BigInt;
+
 use pyo3::exceptions::PyKeyError;
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyString};
 use pyo3::{intern, FromPyObject};
 
+use crate::input::Int;
 use jiter::{cached_py_string, pystring_fast_new, StringCacheMode};
 
 pub trait SchemaDict<'py> {
@@ -133,10 +136,15 @@ pub fn extract_i64(v: &Bound<'_, PyAny>) -> Option<i64> {
         // Can remove this after PyPy 7.3.17 is released
         return None;
     }
-
     v.extract().ok()
 }
 
+pub fn extract_int(v: &Bound<'_, PyAny>) -> Option<Int> {
+    extract_i64(v)
+        .map(Int::I64)
+        .or_else(|| v.extract::<BigInt>().ok().map(Int::Big))
+}
+
 pub(crate) fn new_py_string<'py>(py: Python<'py>, s: &str, cache_str: StringCacheMode) -> Bound<'py, PyString> {
     // we could use `bytecount::num_chars(s.as_bytes()) == s.len()` as orjson does, but it doesn't appear to be faster
     let ascii_only = false;