scirs2_numpy/
extended_int.rs

1//! Extended integer type support: `i128` and `u128`.
2//!
3//! NumPy does not expose native `int128` / `uint128` dtypes, but PyO3 can
4//! round-trip Rust `i128` / `u128` values through Python's arbitrary-precision
5//! `int` type.  This module provides:
6//!
7//! - Free functions [`to_i128`], [`from_i128`], [`to_u128`], [`from_u128`] for
8//!   single-value conversion.
9//! - [`I128Array`] and [`U128Array`] — Python-visible fixed-size containers
10//!   backed by `Vec<i128>` / `Vec<u128>` respectively.
11
12use pyo3::prelude::*;
13use pyo3::types::PyAnyMethods;
14
15// ──────────────────────────────────────────────────────────────────────────────
16// Free-standing conversion helpers
17// ──────────────────────────────────────────────────────────────────────────────
18
19/// Convert a Python `int` to a Rust `i128`.
20///
21/// PyO3 supports extracting `i128` directly from Python's arbitrary-precision
22/// integer type; values outside `[i128::MIN, i128::MAX]` raise `OverflowError`.
23///
24/// # Errors
25/// Returns a [`PyErr`] if the Python object is not an integer or if it
26/// overflows `i128`.
27#[pyfunction]
28pub fn to_i128(val: &Bound<'_, PyAny>) -> PyResult<i128> {
29    val.extract::<i128>()
30}
31
32/// Convert a Rust `i128` to a Python `int`.
33///
34/// The resulting Python object is an ordinary `int` and can be used with any
35/// Python code that accepts integers.
36///
37/// # Errors
38/// Returns a [`PyErr`] on internal conversion failure (should not happen in
39/// practice).
40#[pyfunction]
41pub fn from_i128(py: Python<'_>, val: i128) -> PyResult<Py<PyAny>> {
42    Ok(val.into_pyobject(py)?.into_any().unbind())
43}
44
45/// Convert a Python `int` to a Rust `u128`.
46///
47/// Values outside `[0, u128::MAX]` raise `OverflowError`.
48///
49/// # Errors
50/// Returns a [`PyErr`] if the value is not a non-negative integer or if it
51/// overflows `u128`.
52#[pyfunction]
53pub fn to_u128(val: &Bound<'_, PyAny>) -> PyResult<u128> {
54    val.extract::<u128>()
55}
56
57/// Convert a Rust `u128` to a Python `int`.
58///
59/// # Errors
60/// Returns a [`PyErr`] on internal conversion failure (should not happen in
61/// practice).
62#[pyfunction]
63pub fn from_u128(py: Python<'_>, val: u128) -> PyResult<Py<PyAny>> {
64    Ok(val.into_pyobject(py)?.into_any().unbind())
65}
66
67// ──────────────────────────────────────────────────────────────────────────────
68// I128Array
69// ──────────────────────────────────────────────────────────────────────────────
70
71/// A Python-visible fixed-size array of `i128` values.
72///
73/// NumPy does not expose a native `int128` dtype, so this type provides a
74/// pure-Rust container that can be passed to and from Python via PyO3's
75/// arbitrary-precision integer support.
76#[pyclass(name = "I128Array")]
77pub struct I128Array {
78    /// Backing storage.
79    data: Vec<i128>,
80}
81
82#[pymethods]
83impl I128Array {
84    /// Construct a zero-initialised array of `size` elements.
85    #[new]
86    pub fn new(size: usize) -> Self {
87        Self {
88            data: vec![0i128; size],
89        }
90    }
91
92    /// Return the element at `idx`.
93    ///
94    /// # Errors
95    /// Returns `IndexError` if `idx >= self.len()`.
96    pub fn get(&self, idx: usize) -> PyResult<i128> {
97        self.data.get(idx).copied().ok_or_else(|| {
98            pyo3::exceptions::PyIndexError::new_err(format!(
99                "index {} out of bounds for I128Array of length {}",
100                idx,
101                self.data.len()
102            ))
103        })
104    }
105
106    /// Set the element at `idx` to `val`.
107    ///
108    /// # Errors
109    /// Returns `IndexError` if `idx >= self.len()`.
110    pub fn set(&mut self, idx: usize, val: i128) -> PyResult<()> {
111        let len = self.data.len();
112        self.data.get_mut(idx).map(|e| *e = val).ok_or_else(|| {
113            pyo3::exceptions::PyIndexError::new_err(format!(
114                "index {} out of bounds for I128Array of length {}",
115                idx, len
116            ))
117        })
118    }
119
120    /// Number of elements.
121    pub fn len(&self) -> usize {
122        self.data.len()
123    }
124
125    /// Return `true` if the array is empty.
126    pub fn is_empty(&self) -> bool {
127        self.data.is_empty()
128    }
129
130    /// Compute the sum of all elements.
131    pub fn sum(&self) -> i128 {
132        self.data.iter().sum()
133    }
134
135    /// Return all elements as a Python list of `int` values.
136    pub fn to_list(&self) -> Vec<i128> {
137        self.data.clone()
138    }
139
140    /// Fill the array from a Python list of integers.
141    ///
142    /// # Errors
143    /// Returns a `ValueError` if `values` has a different length than `self`,
144    /// or a `TypeError`/`OverflowError` if an element cannot be converted to
145    /// `i128`.
146    pub fn from_list(&mut self, values: Vec<i128>) -> PyResult<()> {
147        if values.len() != self.data.len() {
148            return Err(pyo3::exceptions::PyValueError::new_err(format!(
149                "expected {} elements, got {}",
150                self.data.len(),
151                values.len()
152            )));
153        }
154        self.data.copy_from_slice(&values);
155        Ok(())
156    }
157}
158
159// ──────────────────────────────────────────────────────────────────────────────
160// U128Array
161// ──────────────────────────────────────────────────────────────────────────────
162
163/// A Python-visible fixed-size array of `u128` values.
164///
165/// Mirrors [`I128Array`] but for unsigned 128-bit integers.
166#[pyclass(name = "U128Array")]
167pub struct U128Array {
168    /// Backing storage.
169    data: Vec<u128>,
170}
171
172#[pymethods]
173impl U128Array {
174    /// Construct a zero-initialised array of `size` elements.
175    #[new]
176    pub fn new(size: usize) -> Self {
177        Self {
178            data: vec![0u128; size],
179        }
180    }
181
182    /// Return the element at `idx`.
183    ///
184    /// # Errors
185    /// Returns `IndexError` if `idx >= self.len()`.
186    pub fn get(&self, idx: usize) -> PyResult<u128> {
187        self.data.get(idx).copied().ok_or_else(|| {
188            pyo3::exceptions::PyIndexError::new_err(format!(
189                "index {} out of bounds for U128Array of length {}",
190                idx,
191                self.data.len()
192            ))
193        })
194    }
195
196    /// Set the element at `idx` to `val`.
197    ///
198    /// # Errors
199    /// Returns `IndexError` if `idx >= self.len()`.
200    pub fn set(&mut self, idx: usize, val: u128) -> PyResult<()> {
201        let len = self.data.len();
202        self.data.get_mut(idx).map(|e| *e = val).ok_or_else(|| {
203            pyo3::exceptions::PyIndexError::new_err(format!(
204                "index {} out of bounds for U128Array of length {}",
205                idx, len
206            ))
207        })
208    }
209
210    /// Number of elements.
211    pub fn len(&self) -> usize {
212        self.data.len()
213    }
214
215    /// Return `true` if the array is empty.
216    pub fn is_empty(&self) -> bool {
217        self.data.is_empty()
218    }
219
220    /// Compute the sum of all elements.
221    pub fn sum(&self) -> u128 {
222        self.data.iter().sum()
223    }
224
225    /// Return all elements as a Python list of `int` values.
226    pub fn to_list(&self) -> Vec<u128> {
227        self.data.clone()
228    }
229
230    /// Fill the array from a Python list of integers.
231    ///
232    /// # Errors
233    /// Returns a `ValueError` if `values` has a different length than `self`,
234    /// or a `TypeError`/`OverflowError` if an element cannot be converted to
235    /// `u128`.
236    pub fn from_list(&mut self, values: Vec<u128>) -> PyResult<()> {
237        if values.len() != self.data.len() {
238            return Err(pyo3::exceptions::PyValueError::new_err(format!(
239                "expected {} elements, got {}",
240                self.data.len(),
241                values.len()
242            )));
243        }
244        self.data.copy_from_slice(&values);
245        Ok(())
246    }
247}
248
249// ──────────────────────────────────────────────────────────────────────────────
250// Module registration
251// ──────────────────────────────────────────────────────────────────────────────
252
253/// Register extended-integer types and conversion functions into a PyO3 module.
254///
255/// Exposes [`I128Array`], [`U128Array`], [`to_i128`], [`from_i128`],
256/// [`to_u128`], and [`from_u128`].
257pub fn register_extended_int_module(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
258    m.add_class::<I128Array>()?;
259    m.add_class::<U128Array>()?;
260    m.add_function(wrap_pyfunction!(to_i128, m)?)?;
261    m.add_function(wrap_pyfunction!(from_i128, m)?)?;
262    m.add_function(wrap_pyfunction!(to_u128, m)?)?;
263    m.add_function(wrap_pyfunction!(from_u128, m)?)?;
264    Ok(())
265}
266
267// ──────────────────────────────────────────────────────────────────────────────
268// Tests
269// ──────────────────────────────────────────────────────────────────────────────
270
271#[cfg(test)]
272mod tests {
273    use super::*;
274
275    #[test]
276    fn i128_array_basic_operations() {
277        let mut arr = I128Array::new(3);
278        assert_eq!(arr.len(), 3);
279        assert_eq!(arr.sum(), 0);
280
281        arr.set(0, 100).expect("set index 0");
282        arr.set(1, 200).expect("set index 1");
283        arr.set(2, 300).expect("set index 2");
284
285        assert_eq!(arr.get(1).expect("get index 1"), 200);
286        assert_eq!(arr.sum(), 600);
287        assert_eq!(arr.to_list(), vec![100i128, 200, 300]);
288    }
289
290    #[test]
291    fn i128_array_out_of_bounds() {
292        let arr = I128Array::new(2);
293        assert!(arr.get(99).is_err());
294    }
295
296    #[test]
297    fn i128_to_from_roundtrip() {
298        Python::attach(|py| {
299            // i128::MAX round-trips through Python int.
300            let py_int = from_i128(py, i128::MAX).expect("from_i128 failed");
301            let val = to_i128(py_int.bind(py)).expect("to_i128 failed");
302            assert_eq!(val, i128::MAX);
303        });
304    }
305
306    #[test]
307    fn u128_roundtrip() {
308        Python::attach(|py| {
309            let py_int = from_u128(py, u128::MAX).expect("from_u128 failed");
310            let val = to_u128(py_int.bind(py)).expect("to_u128 failed");
311            assert_eq!(val, u128::MAX);
312        });
313    }
314
315    #[test]
316    fn u128_array_basic_operations() {
317        let mut arr = U128Array::new(2);
318        arr.set(0, u128::MAX).expect("set 0");
319        assert_eq!(arr.get(0).expect("get 0"), u128::MAX);
320        assert!(arr.get(5).is_err());
321    }
322}
scirs2_numpy/extended_int.rs

scirs2_numpy/
extended_int.rs
