add extract_i64

Author: samuelcolvin

src/errors/location.rs

@@ -8,6 +8,7 @@ use serde::ser::SerializeSeq;
 use serde::{Serialize, Serializer};
 
 use crate::lookup_key::{LookupPath, PathItem};
+use crate::tools::extract_i64;
 
 /// Used to store individual items of the error location, e.g. a string for key/field names
 /// or a number for array indices.
@@ -88,7 +89,7 @@ impl TryFrom<&PyAny> for LocItem {
         if let Ok(py_str) = loc_item.downcast::<PyString>() {
             let str = py_str.to_str()?.to_string();
             Ok(Self::S(str))
-        } else if let Ok(int) = loc_item.extract::<i64>() {
+        } else if let Ok(int) = extract_i64(loc_item) {
             Ok(Self::I(int))
         } else {
             Err(PyTypeError::new_err("Item in a location must be a string or int"))

src/errors/types.rs

@@ -7,7 +7,7 @@ use pyo3::once_cell::GILOnceCell;
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyList};
 
-use crate::tools::{py_err, py_error_type};
+use crate::tools::{extract_i64, py_err, py_error_type};
 use strum::{Display, EnumMessage, IntoEnumIterator};
 use strum_macros::EnumIter;
 
@@ -729,7 +729,7 @@ impl From<String> for Number {
 
 impl FromPyObject<'_> for Number {
     fn extract(obj: &PyAny) -> PyResult<Self> {
-        if let Ok(int) = obj.extract::<i64>() {
+        if let Ok(int) = extract_i64(obj) {
             Ok(Number::Int(int))
         } else if let Ok(float) = obj.extract::<f64>() {
             Ok(Number::Float(float))

src/errors/value_exception.rs

@@ -4,6 +4,7 @@ use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyString};
 
 use crate::input::Input;
+use crate::tools::extract_i64;
 
 use super::{ErrorType, ValError};
 
@@ -74,7 +75,7 @@ impl PydanticCustomError {
                 let key: &PyString = key.downcast()?;
                 if let Ok(py_str) = value.downcast::<PyString>() {
                     message = message.replace(&format!("{{{}}}", key.to_str()?), py_str.to_str()?);
-                } else if let Ok(value_int) = value.extract::<i64>() {
+                } else if let Ok(value_int) = extract_i64(value) {
                     message = message.replace(&format!("{{{}}}", key.to_str()?), &value_int.to_string());
                 } else {
                     // fallback for anything else just in case

src/input/input_python.rs

@@ -11,7 +11,7 @@ use pyo3::types::{PyDictItems, PyDictKeys, PyDictValues};
 use pyo3::{ffi, intern, AsPyPointer, PyTypeInfo};
 
 use crate::errors::{ErrorType, InputValue, LocItem, ValError, ValResult};
-use crate::tools::safe_repr;
+use crate::tools::{extract_i64, safe_repr};
 use crate::{ArgsKwargs, PyMultiHostUrl, PyUrl};
 
 use super::datetime::{
@@ -89,7 +89,7 @@ impl<'a> Input<'a> for PyAny {
     fn as_loc_item(&self) -> LocItem {
         if let Ok(py_str) = self.downcast::<PyString>() {
             py_str.to_string_lossy().as_ref().into()
-        } else if let Ok(key_int) = self.extract::<i64>() {
+        } else if let Ok(key_int) = extract_i64(self) {
             key_int.into()
         } else {
             safe_repr(self).to_string().into()
@@ -256,7 +256,7 @@ impl<'a> Input<'a> for PyAny {
             Ok(bool)
         } else if let Some(cow_str) = maybe_as_string(self, ErrorType::BoolParsing)? {
             str_as_bool(self, &cow_str)
-        } else if let Ok(int) = self.extract::<i64>() {
+        } else if let Ok(int) = extract_i64(self) {
             int_as_bool(self, int)
         } else if let Ok(float) = self.extract::<f64>() {
             match float_as_int(self, float) {
@@ -547,7 +547,7 @@ impl<'a> Input<'a> for PyAny {
             bytes_as_time(self, py_bytes.as_bytes())
         } else if PyBool::is_exact_type_of(self) {
             Err(ValError::new(ErrorType::TimeType, self))
-        } else if let Ok(int) = self.extract::<i64>() {
+        } else if let Ok(int) = extract_i64(self) {
             int_as_time(self, int, 0)
         } else if let Ok(float) = self.extract::<f64>() {
             float_as_time(self, float)
@@ -574,7 +574,7 @@ impl<'a> Input<'a> for PyAny {
             bytes_as_datetime(self, py_bytes.as_bytes())
         } else if PyBool::is_exact_type_of(self) {
             Err(ValError::new(ErrorType::DatetimeType, self))
-        } else if let Ok(int) = self.extract::<i64>() {
+        } else if let Ok(int) = extract_i64(self) {
             int_as_datetime(self, int, 0)
         } else if let Ok(float) = self.extract::<f64>() {
             float_as_datetime(self, float)
@@ -601,7 +601,7 @@ impl<'a> Input<'a> for PyAny {
             bytes_as_timedelta(self, str.as_bytes())
         } else if let Ok(py_bytes) = self.downcast::<PyBytes>() {
             bytes_as_timedelta(self, py_bytes.as_bytes())
-        } else if let Ok(int) = self.extract::<i64>() {
+        } else if let Ok(int) = extract_i64(self) {
             Ok(int_as_duration(self, int)?.into())
         } else if let Ok(float) = self.extract::<f64>() {
             Ok(float_as_duration(self, float)?.into())

src/lookup_key.rs

@@ -8,7 +8,7 @@ use pyo3::types::{PyDict, PyList, PyMapping, PyString};
 use crate::build_tools::py_schema_err;
 use crate::errors::{ErrorType, ValLineError};
 use crate::input::{Input, JsonInput, JsonObject};
-use crate::tools::py_err;
+use crate::tools::{extract_i64, py_err};
 
 /// Used got getting items from python dicts, python objects, or JSON objects, in different ways
 #[derive(Debug, Clone)]
@@ -409,7 +409,7 @@ impl PathItem {
             } else {
                 Ok(Self::Pos(usize_key))
             }
-        } else if let Ok(int_key) = obj.extract::<i64>() {
+        } else if let Ok(int_key) = extract_i64(obj) {
             if index == 0 {
                 py_err!(PyTypeError; "The first item in an alias path should be a string")
             } else {

src/serializers/infer.rs

@@ -14,7 +14,7 @@ use crate::serializers::errors::SERIALIZATION_ERR_MARKER;
 use crate::serializers::filter::SchemaFilter;
 use crate::serializers::shared::{PydanticSerializer, TypeSerializer};
 use crate::serializers::SchemaSerializer;
-use crate::tools::{py_err, safe_repr};
+use crate::tools::{extract_i64, py_err, safe_repr};
 use crate::url::{PyMultiHostUrl, PyUrl};
 
 use super::errors::{py_err_se_err, PydanticSerializationError};
@@ -122,7 +122,7 @@ pub(crate) fn infer_to_python_known(
             // `bool` and `None` can't be subclasses, `ObType::Int`, `ObType::Float`, `ObType::Str` refer to exact types
             ObType::None | ObType::Bool | ObType::Int | ObType::Float | ObType::Str => value.into_py(py),
             // have to do this to make sure subclasses of for example str are upcast to `str`
-            ObType::IntSubclass => value.extract::<i64>()?.into_py(py),
+            ObType::IntSubclass => extract_i64(value)?.into_py(py),
             ObType::FloatSubclass => value.extract::<f64>()?.into_py(py),
             ObType::Decimal => value.to_string().into_py(py),
             ObType::StrSubclass => value.extract::<&str>()?.into_py(py),

src/serializers/type_serializers/literal.rs

@@ -9,7 +9,7 @@ use serde::Serialize;
 
 use crate::build_tools::py_schema_err;
 use crate::definitions::DefinitionsBuilder;
-use crate::tools::SchemaDict;
+use crate::tools::{extract_i64, SchemaDict};
 
 use super::{
     infer_json_key, infer_serialize, infer_to_python, py_err_se_err, BuildSerializer, CombinedSerializer, Extra,
@@ -44,7 +44,7 @@ impl BuildSerializer for LiteralSerializer {
         let mut repr_args: Vec<String> = Vec::new();
         for item in expected {
             repr_args.push(item.repr()?.extract()?);
-            if let Ok(int) = item.extract::<i64>() {
+            if let Ok(int) = extract_i64(item) {
                 expected_int.insert(int);
             } else if let Ok(py_str) = item.downcast::<PyString>() {
                 expected_str.insert(py_str.to_str()?.to_string());
@@ -77,7 +77,7 @@ impl LiteralSerializer {
     fn check<'a>(&self, value: &'a PyAny, extra: &Extra) -> PyResult<OutputValue<'a>> {
         if extra.check.enabled() {
             if !self.expected_int.is_empty() && value.extract::<bool>().is_err() {
-                if let Ok(int) = value.extract::<i64>() {
+                if let Ok(int) = extract_i64(value) {
                     if self.expected_int.contains(&int) {
                         return Ok(OutputValue::OkInt(int));
                     }

src/tools.rs

@@ -1,9 +1,9 @@
 use std::borrow::Cow;
 
-use pyo3::exceptions::PyKeyError;
+use pyo3::exceptions::{PyKeyError, PyTypeError};
 use pyo3::prelude::*;
-use pyo3::types::{PyDict, PyString};
-use pyo3::{intern, FromPyObject};
+use pyo3::types::{PyDict, PyInt, PyString};
+use pyo3::{intern, FromPyObject, PyTypeInfo};
 
 pub trait SchemaDict<'py> {
     fn get_as<T>(&'py self, key: &PyString) -> PyResult<Option<T>>
@@ -98,3 +98,11 @@ pub fn safe_repr(v: &PyAny) -> Cow<str> {
         "<unprintable object>".into()
     }
 }
+
+pub fn extract_i64(v: &PyAny) -> PyResult<i64> {
+    if PyInt::is_type_of(v) {
+        v.extract()
+    } else {
+        py_err!(PyTypeError; "expected int, got {}", safe_repr(v))
+    }
+}

src/validators/union.rs

@@ -15,7 +15,7 @@ use crate::errors::{ErrorType, LocItem, ValError, ValLineError, ValResult};
 use crate::input::{GenericMapping, Input};
 use crate::lookup_key::LookupKey;
 use crate::recursion_guard::RecursionGuard;
-use crate::tools::{py_err, SchemaDict};
+use crate::tools::{extract_i64, py_err, SchemaDict};
 
 use super::custom_error::CustomError;
 use super::{build_validator, BuildValidator, CombinedValidator, Definitions, DefinitionsBuilder, Extra, Validator};
@@ -228,7 +228,7 @@ enum ChoiceKey {
 
 impl ChoiceKey {
     fn from_py(raw: &PyAny) -> PyResult<Self> {
-        if let Ok(py_int) = raw.extract::<i64>() {
+        if let Ok(py_int) = extract_i64(raw) {
             Ok(Self::Int(py_int))
         } else if let Ok(py_str) = raw.downcast::<PyString>() {
             Ok(Self::Str(py_str.to_str()?.to_string()))