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

pyo3/conversions/
num_bigint.rs
