pyo3/conversions/
num_bigint.rs

1#![cfg(feature = "num-bigint")]
2//!  Conversions to and from [num-bigint](https://docs.rs/num-bigint)’s [`BigInt`] and [`BigUint`] types.
3//!
4//! This is useful for converting Python integers when they may not fit in Rust's built-in integer types.
5//!
6//! # Setup
7//!
8//! To use this feature, add this to your **`Cargo.toml`**:
9//!
10//! ```toml
11//! [dependencies]
12//! num-bigint = "*"
13#![doc = concat!("pyo3 = { version = \"", env!("CARGO_PKG_VERSION"),  "\", features = [\"num-bigint\"] }")]
14//! ```
15//!
16//! Note that you must use compatible versions of num-bigint and PyO3.
17//! The required num-bigint version may vary based on the version of PyO3.
18//!
19//! ## Examples
20//!
21//! Using [`BigInt`] to correctly increment an arbitrary precision integer.
22//! This is not possible with Rust's native integers if the Python integer is too large,
23//! in which case it will fail its conversion and raise `OverflowError`.
24//! ```rust,no_run
25//! use num_bigint::BigInt;
26//! use pyo3::prelude::*;
27//!
28//! #[pyfunction]
29//! fn add_one(n: BigInt) -> BigInt {
30//!     n + 1
31//! }
32//!
33//! #[pymodule]
34//! fn my_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
35//!     m.add_function(wrap_pyfunction!(add_one, m)?)?;
36//!     Ok(())
37//! }
38//! ```
39//!
40//! Python code:
41//! ```python
42//! from my_module import add_one
43//!
44//! n = 1 << 1337
45//! value = add_one(n)
46//!
47//! assert n + 1 == value
48//! ```
49
50#[cfg(Py_LIMITED_API)]
51use crate::types::{bytes::PyBytesMethods, PyBytes};
52use crate::{
53    conversion::IntoPyObject, std::num::nb_index, types::PyInt, Borrowed, Bound, FromPyObject,
54    PyAny, PyErr, PyResult, Python,
55};
56
57use num_bigint::{BigInt, BigUint};
58
59#[cfg(not(Py_LIMITED_API))]
60use num_bigint::Sign;
61
62// for identical functionality between BigInt and BigUint
63macro_rules! bigint_conversion {
64    ($rust_ty: ty, $is_signed: literal) => {
65        #[cfg_attr(docsrs, doc(cfg(feature = "num-bigint")))]
66        impl<'py> IntoPyObject<'py> for $rust_ty {
67            type Target = PyInt;
68            type Output = Bound<'py, Self::Target>;
69            type Error = PyErr;
70
71            #[inline]
72            fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
73                (&self).into_pyobject(py)
74            }
75        }
76
77        #[cfg_attr(docsrs, doc(cfg(feature = "num-bigint")))]
78        impl<'py> IntoPyObject<'py> for &$rust_ty {
79            type Target = PyInt;
80            type Output = Bound<'py, Self::Target>;
81            type Error = PyErr;
82
83            fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
84                use num_traits::ToBytes;
85
86                #[cfg(all(not(Py_LIMITED_API), Py_3_13))]
87                {
88                    use crate::conversions::std::num::int_from_ne_bytes;
89                    let bytes = self.to_ne_bytes();
90                    Ok(int_from_ne_bytes::<{ $is_signed }>(py, &bytes))
91                }
92
93                #[cfg(all(not(Py_LIMITED_API), not(Py_3_13)))]
94                {
95                    use crate::conversions::std::num::int_from_le_bytes;
96                    let bytes = self.to_le_bytes();
97                    Ok(int_from_le_bytes::<{ $is_signed }>(py, &bytes))
98                }
99
100                #[cfg(Py_LIMITED_API)]
101                {
102                    use $crate::py_result_ext::PyResultExt;
103                    use $crate::types::any::PyAnyMethods;
104                    let bytes = self.to_le_bytes();
105                    let bytes_obj = PyBytes::new(py, &bytes);
106                    let kwargs = if $is_signed {
107                        let kwargs = crate::types::PyDict::new(py);
108                        kwargs.set_item(crate::intern!(py, "signed"), true)?;
109                        Some(kwargs)
110                    } else {
111                        None
112                    };
113                    unsafe {
114                        py.get_type::<PyInt>()
115                            .call_method("from_bytes", (bytes_obj, "little"), kwargs.as_ref())
116                            .cast_into_unchecked()
117                    }
118                }
119            }
120        }
121    };
122}
123
124bigint_conversion!(BigUint, false);
125bigint_conversion!(BigInt, true);
126
127#[cfg_attr(docsrs, doc(cfg(feature = "num-bigint")))]
128impl<'py> FromPyObject<'_, 'py> for BigInt {
129    type Error = PyErr;
130
131    fn extract(ob: Borrowed<'_, 'py, PyAny>) -> Result<BigInt, Self::Error> {
132        // fast path - checking for subclass of `int` just checks a bit in the type object
133        let num_owned: Bound<'_, PyInt>;
134        let num = if let Ok(long) = ob.cast::<PyInt>() {
135            long
136        } else {
137            num_owned = nb_index(&ob)?;
138            num_owned.as_borrowed()
139        };
140        #[cfg(not(Py_LIMITED_API))]
141        {
142            let mut buffer = int_to_u32_vec::<true>(&num)?;
143            let sign = if buffer.last().copied().is_some_and(|last| last >> 31 != 0) {
144                // BigInt::new takes an unsigned array, so need to convert from two's complement
145                // flip all bits, 'subtract' 1 (by adding one to the unsigned array)
146                let mut elements = buffer.iter_mut();
147                for element in elements.by_ref() {
148                    *element = (!*element).wrapping_add(1);
149                    if *element != 0 {
150                        // if the element didn't wrap over, no need to keep adding further ...
151                        break;
152                    }
153                }
154                // ... so just two's complement the rest
155                for element in elements {
156                    *element = !*element;
157                }
158                Sign::Minus
159            } else {
160                Sign::Plus
161            };
162            Ok(BigInt::new(sign, buffer))
163        }
164        #[cfg(Py_LIMITED_API)]
165        {
166            let n_bits = int_n_bits(&num)?;
167            if n_bits == 0 {
168                return Ok(BigInt::from(0isize));
169            }
170            let bytes = int_to_py_bytes(&num, (n_bits + 8) / 8, true)?;
171            Ok(BigInt::from_signed_bytes_le(bytes.as_bytes()))
172        }
173    }
174}
175
176#[cfg_attr(docsrs, doc(cfg(feature = "num-bigint")))]
177impl<'py> FromPyObject<'_, 'py> for BigUint {
178    type Error = PyErr;
179
180    fn extract(ob: Borrowed<'_, 'py, PyAny>) -> Result<BigUint, Self::Error> {
181        // fast path - checking for subclass of `int` just checks a bit in the type object
182        let num_owned: Bound<'_, PyInt>;
183        let num = if let Ok(long) = ob.cast::<PyInt>() {
184            long
185        } else {
186            num_owned = nb_index(&ob)?;
187            num_owned.as_borrowed()
188        };
189        #[cfg(not(Py_LIMITED_API))]
190        {
191            let buffer = int_to_u32_vec::<false>(&num)?;
192            Ok(BigUint::new(buffer))
193        }
194        #[cfg(Py_LIMITED_API)]
195        {
196            let n_bits = int_n_bits(&num)?;
197            if n_bits == 0 {
198                return Ok(BigUint::from(0usize));
199            }
200            let bytes = int_to_py_bytes(&num, n_bits.div_ceil(8), false)?;
201            Ok(BigUint::from_bytes_le(bytes.as_bytes()))
202        }
203    }
204}
205
206#[cfg(not(any(Py_LIMITED_API, Py_3_13)))]
207#[inline]
208fn int_to_u32_vec<const SIGNED: bool>(long: &Bound<'_, PyInt>) -> PyResult<Vec<u32>> {
209    use crate::ffi;
210
211    let mut buffer = Vec::new();
212    let n_bits = int_n_bits(long)?;
213    if n_bits == 0 {
214        return Ok(buffer);
215    }
216    let n_digits = if SIGNED {
217        (n_bits + 32) / 32
218    } else {
219        n_bits.div_ceil(32)
220    };
221    buffer.reserve_exact(n_digits);
222    unsafe {
223        crate::err::error_on_minusone(
224            long.py(),
225            ffi::_PyLong_AsByteArray(
226                long.as_ptr().cast(),
227                buffer.as_mut_ptr() as *mut u8,
228                n_digits * 4,
229                1,
230                SIGNED.into(),
231            ),
232        )?;
233        buffer.set_len(n_digits)
234    };
235    buffer
236        .iter_mut()
237        .for_each(|chunk| *chunk = u32::from_le(*chunk));
238
239    Ok(buffer)
240}
241
242#[cfg(all(not(Py_LIMITED_API), Py_3_13))]
243#[inline]
244fn int_to_u32_vec<const SIGNED: bool>(long: &Bound<'_, PyInt>) -> PyResult<Vec<u32>> {
245    use crate::ffi;
246
247    let mut buffer = Vec::new();
248    let mut flags = ffi::Py_ASNATIVEBYTES_LITTLE_ENDIAN;
249    if !SIGNED {
250        flags |= ffi::Py_ASNATIVEBYTES_UNSIGNED_BUFFER | ffi::Py_ASNATIVEBYTES_REJECT_NEGATIVE;
251    }
252    let n_bytes =
253        unsafe { ffi::PyLong_AsNativeBytes(long.as_ptr().cast(), std::ptr::null_mut(), 0, flags) };
254    let n_bytes_unsigned: usize = n_bytes
255        .try_into()
256        .map_err(|_| crate::PyErr::fetch(long.py()))?;
257    if n_bytes == 0 {
258        return Ok(buffer);
259    }
260    let n_digits = n_bytes_unsigned.div_ceil(4);
261    buffer.reserve_exact(n_digits);
262    unsafe {
263        ffi::PyLong_AsNativeBytes(
264            long.as_ptr().cast(),
265            buffer.as_mut_ptr().cast(),
266            (n_digits * 4).try_into().unwrap(),
267            flags,
268        );
269        buffer.set_len(n_digits);
270    };
271    buffer
272        .iter_mut()
273        .for_each(|chunk| *chunk = u32::from_le(*chunk));
274
275    Ok(buffer)
276}
277
278#[cfg(Py_LIMITED_API)]
279fn int_to_py_bytes<'py>(
280    long: &Bound<'py, PyInt>,
281    n_bytes: usize,
282    is_signed: bool,
283) -> PyResult<Bound<'py, PyBytes>> {
284    use crate::intern;
285    use crate::types::any::PyAnyMethods;
286    let py = long.py();
287    let kwargs = if is_signed {
288        let kwargs = crate::types::PyDict::new(py);
289        kwargs.set_item(intern!(py, "signed"), true)?;
290        Some(kwargs)
291    } else {
292        None
293    };
294    let bytes = long.call_method(
295        intern!(py, "to_bytes"),
296        (n_bytes, intern!(py, "little")),
297        kwargs.as_ref(),
298    )?;
299    Ok(bytes.cast_into()?)
300}
301
302#[inline]
303#[cfg(any(not(Py_3_13), Py_LIMITED_API))]
304fn int_n_bits(long: &Bound<'_, PyInt>) -> PyResult<usize> {
305    let py = long.py();
306    #[cfg(not(Py_LIMITED_API))]
307    {
308        // fast path
309        let n_bits = unsafe { crate::ffi::_PyLong_NumBits(long.as_ptr()) };
310        if n_bits == (-1isize as usize) {
311            return Err(crate::PyErr::fetch(py));
312        }
313        Ok(n_bits)
314    }
315
316    #[cfg(Py_LIMITED_API)]
317    {
318        // slow path
319        use crate::types::PyAnyMethods;
320        long.call_method0(crate::intern!(py, "bit_length"))
321            .and_then(|any| any.extract())
322    }
323}
324
325#[cfg(test)]
326mod tests {
327    use super::*;
328    use crate::exceptions::PyTypeError;
329    use crate::test_utils::generate_unique_module_name;
330    use crate::types::{PyAnyMethods as _, PyDict, PyModule};
331    use pyo3_ffi::c_str;
332
333    fn rust_fib<T>() -> impl Iterator<Item = T>
334    where
335        T: From<u16>,
336        for<'a> &'a T: std::ops::Add<Output = T>,
337    {
338        let mut f0: T = T::from(1);
339        let mut f1: T = T::from(1);
340        std::iter::from_fn(move || {
341            let f2 = &f0 + &f1;
342            Some(std::mem::replace(&mut f0, std::mem::replace(&mut f1, f2)))
343        })
344    }
345
346    fn python_fib(py: Python<'_>) -> impl Iterator<Item = Bound<'_, PyAny>> + '_ {
347        let mut f0 = 1i32.into_pyobject(py).unwrap().into_any();
348        let mut f1 = 1i32.into_pyobject(py).unwrap().into_any();
349        std::iter::from_fn(move || {
350            let f2 = f0.call_method1("__add__", (&f1,)).unwrap();
351            Some(std::mem::replace(&mut f0, std::mem::replace(&mut f1, f2)))
352        })
353    }
354
355    #[test]
356    fn convert_biguint() {
357        Python::attach(|py| {
358            // check the first 2000 numbers in the fibonacci sequence
359            for (py_result, rs_result) in python_fib(py).zip(rust_fib::<BigUint>()).take(2000) {
360                // Python -> Rust
361                assert_eq!(py_result.extract::<BigUint>().unwrap(), rs_result);
362                // Rust -> Python
363                assert!(py_result.eq(rs_result).unwrap());
364            }
365        });
366    }
367
368    #[test]
369    fn convert_bigint() {
370        Python::attach(|py| {
371            // check the first 2000 numbers in the fibonacci sequence
372            for (py_result, rs_result) in python_fib(py).zip(rust_fib::<BigInt>()).take(2000) {
373                // Python -> Rust
374                assert_eq!(py_result.extract::<BigInt>().unwrap(), rs_result);
375                // Rust -> Python
376                assert!(py_result.eq(&rs_result).unwrap());
377
378                // negate
379
380                let rs_result = rs_result * -1;
381                let py_result = py_result.call_method0("__neg__").unwrap();
382
383                // Python -> Rust
384                assert_eq!(py_result.extract::<BigInt>().unwrap(), rs_result);
385                // Rust -> Python
386                assert!(py_result.eq(rs_result).unwrap());
387            }
388        });
389    }
390
391    fn python_index_class(py: Python<'_>) -> Bound<'_, PyModule> {
392        let index_code = c_str!(
393            r#"
394class C:
395    def __init__(self, x):
396        self.x = x
397    def __index__(self):
398        return self.x
399"#
400        );
401        PyModule::from_code(
402            py,
403            index_code,
404            c"index.py",
405            &generate_unique_module_name("index"),
406        )
407        .unwrap()
408    }
409
410    #[test]
411    fn convert_index_class() {
412        Python::attach(|py| {
413            let index = python_index_class(py);
414            let locals = PyDict::new(py);
415            locals.set_item("index", index).unwrap();
416            let ob = py.eval(c"index.C(10)", None, Some(&locals)).unwrap();
417            let _: BigInt = ob.extract().unwrap();
418            let _: BigUint = ob.extract().unwrap();
419        });
420    }
421
422    #[test]
423    fn handle_zero() {
424        Python::attach(|py| {
425            let zero: BigInt = 0i32.into_pyobject(py).unwrap().extract().unwrap();
426            assert_eq!(zero, BigInt::from(0));
427        })
428    }
429
430    /// `OverflowError` on converting Python int to BigInt, see issue #629
431    #[test]
432    fn check_overflow() {
433        Python::attach(|py| {
434            macro_rules! test {
435                ($T:ty, $value:expr, $py:expr) => {
436                    let value = $value;
437                    println!("{}: {}", stringify!($T), value);
438                    let python_value = value.clone().into_pyobject(py).unwrap();
439                    let roundtrip_value = python_value.extract::<$T>().unwrap();
440                    assert_eq!(value, roundtrip_value);
441                };
442            }
443
444            for i in 0..=256usize {
445                // test a lot of values to help catch other bugs too
446                test!(BigInt, BigInt::from(i), py);
447                test!(BigUint, BigUint::from(i), py);
448                test!(BigInt, -BigInt::from(i), py);
449                test!(BigInt, BigInt::from(1) << i, py);
450                test!(BigUint, BigUint::from(1u32) << i, py);
451                test!(BigInt, -BigInt::from(1) << i, py);
452                test!(BigInt, (BigInt::from(1) << i) + 1u32, py);
453                test!(BigUint, (BigUint::from(1u32) << i) + 1u32, py);
454                test!(BigInt, (-BigInt::from(1) << i) + 1u32, py);
455                test!(BigInt, (BigInt::from(1) << i) - 1u32, py);
456                test!(BigUint, (BigUint::from(1u32) << i) - 1u32, py);
457                test!(BigInt, (-BigInt::from(1) << i) - 1u32, py);
458            }
459        });
460    }
461
462    #[test]
463    fn from_py_float_type_error() {
464        Python::attach(|py| {
465            let obj = (12.3f64).into_pyobject(py).unwrap();
466            let err = obj.extract::<BigInt>().unwrap_err();
467            assert!(err.is_instance_of::<PyTypeError>(py));
468
469            let err = obj.extract::<BigUint>().unwrap_err();
470            assert!(err.is_instance_of::<PyTypeError>(py));
471        });
472    }
473}