num_bigint.rs

// Copyright (c) 2017-present PyO3 Project and Contributors
//
// based on Daniel Grunwald's https://github.com/dgrunwald/rust-cpython

#![cfg(all(feature = "num-bigint", not(any(Py_LIMITED_API, PyPy))))]
#![cfg_attr(
    docsrs,
    doc(cfg(all(feature = "num-bigint", not(any(Py_LIMITED_API, PyPy)))))
)]
//!  Conversions to and from [num-bigint](https://docs.rs/num-bigint)’s [`BigInt`] and [`BigUint`] types.
//!
//! This is useful for converting Python integers when they may not fit in Rust's built-in integer types.
//!
//! # Setup
//!
//! To use this feature, add this to your **`Cargo.toml`**:
//!
//! ```toml
//! [dependencies]
//! num-bigint = "*"
// workaround for `extended_key_value_attributes`: https://github.com/rust-lang/rust/issues/82768#issuecomment-803935643
#![cfg_attr(docsrs, cfg_attr(docsrs, doc = concat!("pyo3 = { version = \"", env!("CARGO_PKG_VERSION"),  "\", features = [\"num-bigint\"] }")))]
#![cfg_attr(
    not(docsrs),
    doc = "pyo3 = { version = \"*\", features = [\"num-bigint\"] }"
)]
//! ```
//!
//! Note that you must use compatible versions of num-bigint and PyO3.
//! The required num-bigint version may vary based on the version of PyO3.
//!
//! ## Examples
//!
//! Using [`BigInt`] to correctly increment an arbitrary precision integer.
//! This is not possible with Rust's native integers if the Python integer is too large,
//! in which case it will fail its conversion and raise `OverflowError`.
//!
//! ```rust
//! use num_bigint::BigInt;
//! use pyo3::prelude::*;
//!
//! #[pyfunction]
//! fn add_one(n: BigInt) -> BigInt {
//!     n + 1
//! }
//!
//! #[pymodule]
//! fn my_module(_py: Python, m: &PyModule) -> PyResult<()> {
//!     m.add_function(wrap_pyfunction!(add_one, m)?)?;
//!     Ok(())
//! }
//! ```
//!
//! Python code:
//! ```python
//! from my_module import add_one
//!
//! n = 1 << 1337
//! value = add_one(n)
//!
//! assert n + 1 == value
//! ```

use crate::{
    err, ffi, types::*, AsPyPointer, FromPyObject, IntoPy, Py, PyAny, PyErr, PyObject, PyResult,
    Python, ToPyObject,
};

use num_bigint::{BigInt, BigUint};
use std::os::raw::{c_int, c_uchar};

#[cfg(not(all(windows, PyPy)))]
unsafe fn extract(ob: &PyLong, buffer: &mut [c_uchar], is_signed: c_int) -> PyResult<()> {
    err::error_on_minusone(
        ob.py(),
        ffi::_PyLong_AsByteArray(
            ob.as_ptr() as *mut ffi::PyLongObject,
            buffer.as_mut_ptr(),
            buffer.len(),
            1,
            is_signed,
        ),
    )
}

macro_rules! bigint_conversion {
    ($rust_ty: ty, $is_signed: expr, $to_bytes: path, $from_bytes: path) => {
        impl ToPyObject for $rust_ty {
            fn to_object(&self, py: Python) -> PyObject {
                unsafe {
                    let bytes = $to_bytes(self);
                    let obj = ffi::_PyLong_FromByteArray(
                        bytes.as_ptr() as *const c_uchar,
                        bytes.len(),
                        1,
                        $is_signed,
                    );
                    PyObject::from_owned_ptr(py, obj)
                }
            }
        }
        impl IntoPy<PyObject> for $rust_ty {
            fn into_py(self, py: Python) -> PyObject {
                self.to_object(py)
            }
        }
        impl<'source> FromPyObject<'source> for $rust_ty {
            fn extract(ob: &'source PyAny) -> PyResult<$rust_ty> {
                let py = ob.py();
                unsafe {
                    let num = ffi::PyNumber_Index(ob.as_ptr());
                    if num.is_null() {
                        return Err(PyErr::api_call_failed(py));
                    }
                    let n_bits = ffi::_PyLong_NumBits(num);
                    let n_bytes = if n_bits < 0 {
                        return Err(PyErr::api_call_failed(py));
                    } else if n_bits == 0 {
                        0
                    } else {
                        (n_bits as usize - 1 + $is_signed) / 8 + 1
                    };
                    let num: Py<PyLong> = Py::from_owned_ptr(py, num);
                    if n_bytes <= 128 {
                        let mut buffer = [0; 128];
                        extract(num.as_ref(py), &mut buffer[..n_bytes], $is_signed)?;
                        Ok($from_bytes(&buffer[..n_bytes]))
                    } else {
                        let mut buffer = vec![0; n_bytes];
                        extract(num.as_ref(py), &mut buffer, $is_signed)?;
                        Ok($from_bytes(&buffer))
                    }
                }
            }
        }
    };
}
bigint_conversion!(BigUint, 0, BigUint::to_bytes_le, BigUint::from_bytes_le);
bigint_conversion!(
    BigInt,
    1,
    BigInt::to_signed_bytes_le,
    BigInt::from_signed_bytes_le
);

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{PyDict, PyModule};
    use indoc::indoc;

    fn python_fib(py: Python) -> &PyModule {
        let fib_code = indoc!(
            r#"
                def fib(n):
                    f0, f1 = 0, 1
                    for _ in range(n):
                        f0, f1 = f1, f0 + f1
                    return f0

                def fib_neg(n):
                    return -fib(n)
        "#
        );
        PyModule::from_code(py, fib_code, "fib.py", "fib").unwrap()
    }

    fn rust_fib<T>(n: usize) -> T
    where
        T: From<u16>,
        for<'a> &'a T: std::ops::Add<Output = T>,
    {
        let mut f0: T = T::from(0);
        let mut f1: T = T::from(1);
        for _ in 0..n {
            let f2 = &f0 + &f1;
            f0 = std::mem::replace(&mut f1, f2);
        }
        f0
    }

    #[test]
    fn convert_biguint() {
        Python::with_gil(|py| {
            let rs_result: BigUint = rust_fib(400);
            let fib = python_fib(py);
            let locals = PyDict::new(py);
            locals.set_item("rs_result", &rs_result).unwrap();
            locals.set_item("fib", fib).unwrap();
            // Checks if Rust BigUint -> Python Long conversion is correct
            py.run("assert fib.fib(400) == rs_result", None, Some(locals))
                .unwrap();
            // Checks if Python Long -> Rust BigUint conversion is correct if N is small
            let py_result: BigUint =
                FromPyObject::extract(fib.getattr("fib").unwrap().call1((400,)).unwrap()).unwrap();
            assert_eq!(rs_result, py_result);
            // Checks if Python Long -> Rust BigUint conversion is correct if N is large
            let rs_result: BigUint = rust_fib(2000);
            let py_result: BigUint =
                FromPyObject::extract(fib.getattr("fib").unwrap().call1((2000,)).unwrap()).unwrap();
            assert_eq!(rs_result, py_result);
        });
    }

    #[test]
    fn convert_bigint() {
        Python::with_gil(|py| {
            let rs_result = rust_fib::<BigInt>(400) * -1;
            let fib = python_fib(py);
            let locals = PyDict::new(py);
            locals.set_item("rs_result", &rs_result).unwrap();
            locals.set_item("fib", fib).unwrap();
            // Checks if Rust BigInt -> Python Long conversion is correct
            py.run("assert fib.fib_neg(400) == rs_result", None, Some(locals))
                .unwrap();
            // Checks if Python Long -> Rust BigInt conversion is correct if N is small
            let py_result: BigInt =
                FromPyObject::extract(fib.getattr("fib_neg").unwrap().call1((400,)).unwrap())
                    .unwrap();
            assert_eq!(rs_result, py_result);
            // Checks if Python Long -> Rust BigInt conversion is correct if N is large
            let rs_result = rust_fib::<BigInt>(2000) * -1;
            let py_result: BigInt =
                FromPyObject::extract(fib.getattr("fib_neg").unwrap().call1((2000,)).unwrap())
                    .unwrap();
            assert_eq!(rs_result, py_result);
        })
    }

    fn python_index_class(py: Python) -> &PyModule {
        let index_code = indoc!(
            r#"
                class C:
                    def __init__(self, x):
                        self.x = x
                    def __index__(self):
                        return self.x
                "#
        );
        PyModule::from_code(py, index_code, "index.py", "index").unwrap()
    }

    #[test]
    fn convert_index_class() {
        Python::with_gil(|py| {
            let index = python_index_class(py);
            let locals = PyDict::new(py);
            locals.set_item("index", index).unwrap();
            let ob = py.eval("index.C(10)", None, Some(locals)).unwrap();
            let _: BigInt = FromPyObject::extract(ob).unwrap();
        });
    }

    #[test]
    fn handle_zero() {
        Python::with_gil(|py| {
            let fib = python_fib(py);
            let zero: BigInt =
                FromPyObject::extract(fib.getattr("fib").unwrap().call1((0,)).unwrap()).unwrap();
            assert_eq!(zero, BigInt::from(0));
        })
    }

    /// `OverflowError` on converting Python int to BigInt, see issue #629
    #[test]
    fn check_overflow() {
        Python::with_gil(|py| {
            macro_rules! test {
                ($T:ty, $value:expr, $py:expr) => {
                    let value = $value;
                    println!("{}: {}", stringify!($T), value);
                    let python_value = value.clone().to_object(py);
                    let roundtrip_value = python_value.extract::<$T>(py).unwrap();
                    assert_eq!(value, roundtrip_value);
                };
            }

            for i in 0..=256usize {
                // test a lot of values to help catch other bugs too
                test!(BigInt, BigInt::from(i), py);
                test!(BigUint, BigUint::from(i), py);
                test!(BigInt, -BigInt::from(i), py);
                test!(BigInt, BigInt::from(1) << i, py);
                test!(BigUint, BigUint::from(1u32) << i, py);
                test!(BigInt, -BigInt::from(1) << i, py);
                test!(BigInt, (BigInt::from(1) << i) + 1u32, py);
                test!(BigUint, (BigUint::from(1u32) << i) + 1u32, py);
                test!(BigInt, (-BigInt::from(1) << i) + 1u32, py);
                test!(BigInt, (BigInt::from(1) << i) - 1u32, py);
                test!(BigUint, (BigUint::from(1u32) << i) - 1u32, py);
                test!(BigInt, (-BigInt::from(1) << i) - 1u32, py);
            }
        });
    }
}