Merge pull request #3742 from samuelcolvin/int-extraction-performance

improve performance of successful int extract by ~30% by avoiding calls to `__index__` where redundant

newsfragments/3742.changed.md

@@ -0,0 +1 @@
+Improve performance of `extract::<i64>` (and other integer types) by avoiding call to `__index__()` converting the value to an integer for 3.10+. Gives performance improvement of around 30% for successful extraction.

src/conversions/std/num.rs

@@ -44,8 +44,39 @@ macro_rules! int_fits_larger_int {
     };
 }
 
+macro_rules! extract_int {
+    ($obj:ident, $error_val:expr, $pylong_as:expr) => {
+        extract_int!($obj, $error_val, $pylong_as, false)
+    };
+
+    ($obj:ident, $error_val:expr, $pylong_as:expr, $force_index_call: literal) => {
+        // In python 3.8+ `PyLong_AsLong` and friends takes care of calling `PyNumber_Index`,
+        // however 3.8 & 3.9 do lossy conversion of floats, hence we only use the
+        // simplest logic for 3.10+ where that was fixed - python/cpython#82180.
+        // `PyLong_AsUnsignedLongLong` does not call `PyNumber_Index`, hence the `force_index_call` argument
+        // See https://github.com/PyO3/pyo3/pull/3742 for detials
+        if cfg!(Py_3_10) && !$force_index_call {
+            err_if_invalid_value($obj.py(), $error_val, unsafe { $pylong_as($obj.as_ptr()) })
+        } else if let Ok(long) = $obj.downcast::<crate::types::PyLong>() {
+            // fast path - checking for subclass of `int` just checks a bit in the type $object
+            err_if_invalid_value($obj.py(), $error_val, unsafe { $pylong_as(long.as_ptr()) })
+        } else {
+            unsafe {
+                let num = ffi::PyNumber_Index($obj.as_ptr());
+                if num.is_null() {
+                    Err(PyErr::fetch($obj.py()))
+                } else {
+                    let result = err_if_invalid_value($obj.py(), $error_val, $pylong_as(num));
+                    ffi::Py_DECREF(num);
+                    result
+                }
+            }
+        }
+    };
+}
+
 macro_rules! int_convert_u64_or_i64 {
-    ($rust_type:ty, $pylong_from_ll_or_ull:expr, $pylong_as_ll_or_ull:expr) => {
+    ($rust_type:ty, $pylong_from_ll_or_ull:expr, $pylong_as_ll_or_ull:expr, $force_index_call:literal) => {
         impl ToPyObject for $rust_type {
             #[inline]
             fn to_object(&self, py: Python<'_>) -> PyObject {
@@ -64,18 +95,8 @@ macro_rules! int_convert_u64_or_i64 {
             }
         }
         impl<'source> FromPyObject<'source> for $rust_type {
-            fn extract(ob: &'source PyAny) -> PyResult<$rust_type> {
-                let ptr = ob.as_ptr();
-                unsafe {
-                    let num = ffi::PyNumber_Index(ptr);
-                    if num.is_null() {
-                        Err(PyErr::fetch(ob.py()))
-                    } else {
-                        let result = err_if_invalid_value(ob.py(), !0, $pylong_as_ll_or_ull(num));
-                        ffi::Py_DECREF(num);
-                        result
-                    }
-                }
+            fn extract(obj: &'source PyAny) -> PyResult<$rust_type> {
+                extract_int!(obj, !0, $pylong_as_ll_or_ull, $force_index_call)
             }
 
             #[cfg(feature = "experimental-inspect")]
@@ -106,17 +127,7 @@ macro_rules! int_fits_c_long {
 
         impl<'source> FromPyObject<'source> for $rust_type {
             fn extract(obj: &'source PyAny) -> PyResult<Self> {
-                let ptr = obj.as_ptr();
-                let val = unsafe {
-                    let num = ffi::PyNumber_Index(ptr);
-                    if num.is_null() {
-                        Err(PyErr::fetch(obj.py()))
-                    } else {
-                        let val = err_if_invalid_value(obj.py(), -1, ffi::PyLong_AsLong(num));
-                        ffi::Py_DECREF(num);
-                        val
-                    }
-                }?;
+                let val: c_long = extract_int!(obj, -1, ffi::PyLong_AsLong)?;
                 <$rust_type>::try_from(val)
                     .map_err(|e| exceptions::PyOverflowError::new_err(e.to_string()))
             }
@@ -146,7 +157,7 @@ int_fits_c_long!(i64);
 
 // manual implementation for i64 on systems with 32-bit long
 #[cfg(any(target_pointer_width = "32", target_os = "windows"))]
-int_convert_u64_or_i64!(i64, ffi::PyLong_FromLongLong, ffi::PyLong_AsLongLong);
+int_convert_u64_or_i64!(i64, ffi::PyLong_FromLongLong, ffi::PyLong_AsLongLong, false);
 
 #[cfg(all(target_pointer_width = "64", not(target_os = "windows")))]
 int_fits_c_long!(isize);
@@ -159,7 +170,8 @@ int_fits_larger_int!(usize, u64);
 int_convert_u64_or_i64!(
     u64,
     ffi::PyLong_FromUnsignedLongLong,
-    ffi::PyLong_AsUnsignedLongLong
+    ffi::PyLong_AsUnsignedLongLong,
+    true
 );
 
 #[cfg(not(Py_LIMITED_API))]
@@ -738,4 +750,27 @@ mod tests {
     test_nonzero_common!(nonzero_usize, NonZeroUsize);
     test_nonzero_common!(nonzero_i128, NonZeroI128);
     test_nonzero_common!(nonzero_u128, NonZeroU128);
+
+    #[test]
+    fn test_i64_bool() {
+        Python::with_gil(|py| {
+            let obj = true.to_object(py);
+            assert_eq!(1, obj.extract::<i64>(py).unwrap());
+            let obj = false.to_object(py);
+            assert_eq!(0, obj.extract::<i64>(py).unwrap());
+        })
+    }
+
+    #[test]
+    fn test_i64_f64() {
+        Python::with_gil(|py| {
+            let obj = 12.34f64.to_object(py);
+            let err = obj.extract::<i64>(py).unwrap_err();
+            assert!(err.is_instance_of::<crate::exceptions::PyTypeError>(py));
+            // with no remainder
+            let obj = 12f64.to_object(py);
+            let err = obj.extract::<i64>(py).unwrap_err();
+            assert!(err.is_instance_of::<crate::exceptions::PyTypeError>(py));
+        })
+    }
 }