pyo3 variant working

src/lib.rs

@@ -5,3 +5,159 @@ pub mod record;
 pub mod records;
 pub mod elements;
 pub mod library;
+
+
+//use ndarray;
+use numpy::{PyArray1, PyUntypedArray, PyUntypedArrayMethods, PyArrayDescrMethods, PyArrayMethods, dtype};
+use pyo3::prelude::{Python, pymodule, PyModule, PyResult, Bound, wrap_pyfunction, pyfunction, PyModuleMethods, PyAnyMethods};
+use pyo3::exceptions::{PyValueError, PyTypeError};
+
+
+#[pymodule]
+fn klamath_rs_ext(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_function(wrap_pyfunction!(arr_to_int2, m)?)?;
+    m.add_function(wrap_pyfunction!(arr_to_int4, m)?)?;
+    Ok(())
+}
+
+
+#[pyfunction]
+fn arr_to_int2(py: Python<'_>, pyarr: &Bound<'_, PyUntypedArray>) -> PyResult<()> {
+    use rust_util::ToInt2BE;
+
+    assert!(pyarr.is_c_contiguous(), "Array must be c-contiguous!");
+
+    macro_rules! i2if {
+        ( $el_type:expr, $tt:ty ) => {
+            if $el_type.is_equiv_to(&dtype::<$tt>(py)) {
+                let arr = pyarr.downcast::<PyArray1<$tt>>()?;
+                let mut array = unsafe { arr.as_array_mut() };
+                for xx in array.iter_mut() {
+                    *xx = <$tt>::convert_to_i2be(*xx).map_err(
+                        |e| PyValueError::new_err(format!("Invalid value for 2-byte int: {}", e))
+                        )?;
+                }
+                return Ok(())
+            }
+        }
+    }
+
+
+    let el_type = pyarr.dtype();
+    i2if!(el_type, f64);
+    i2if!(el_type, f32);
+    i2if!(el_type, i64);
+    i2if!(el_type, u64);
+    i2if!(el_type, i32);
+    i2if!(el_type, u32);
+    i2if!(el_type, i16);
+    i2if!(el_type, u16);
+
+    Err(PyTypeError::new_err(format!("arr_to_int2 not implemented for type {:?}", el_type)))
+}
+
+
+#[pyfunction]
+fn arr_to_int4(py: Python<'_>, pyarr: &Bound<'_, PyUntypedArray>) -> PyResult<()> {
+    use rust_util::ToInt4BE;
+
+    assert!(pyarr.is_c_contiguous(), "Array must be c-contiguous!");
+
+    macro_rules! i4if {
+        ( $el_type:expr, $tt:ty ) => {
+            if $el_type.is_equiv_to(&dtype::<$tt>(py)) {
+                let arr = pyarr.downcast::<PyArray1<$tt>>()?;
+                let mut array = unsafe { arr.as_array_mut() };
+                for xx in array.iter_mut() {
+                    *xx = <$tt>::convert_to_i4be(*xx).map_err(
+                        |e| PyValueError::new_err(format!("Invalid value for 4-byte int: {}", e))
+                        )?;
+                }
+                return Ok(())
+            }
+        }
+    }
+
+    let el_type = pyarr.dtype();
+    i4if!(el_type, f64);
+    i4if!(el_type, f32);
+    i4if!(el_type, i64);
+    i4if!(el_type, u64);
+    i4if!(el_type, i32);
+    i4if!(el_type, u32);
+
+    Err(PyTypeError::new_err(format!("arr_to_int4 not implemented for type {:?}", el_type)))
+}
+
+
+mod rust_util {
+    use byteorder::{ByteOrder, BigEndian};
+    use std::mem::size_of;
+
+    pub trait ToInt2BE {
+        fn convert_to_i2be(ii: Self) -> Result<Self, Self> where Self: Sized;
+    }
+
+    pub trait ToInt4BE {
+        fn convert_to_i4be(ii: Self) -> Result<Self, Self> where Self: Sized;
+    }
+
+    macro_rules! impl_i2be {
+        ( $tt:ty ) => {
+            impl ToInt2BE for $tt {
+                fn convert_to_i2be(ii: $tt) -> Result<$tt, $tt> {
+                    if ii < i16::MIN as $tt { return Err(ii); }
+                    if ii > i16::MAX as $tt { return Err(ii); }
+
+                    let mut buf = [0; size_of::<$tt>()];
+                    BigEndian::write_i16(&mut buf, ii as i16);
+                    Ok(<$tt>::from_le_bytes(buf))
+                }
+            }
+        }
+    }
+
+    macro_rules! impl_i4be {
+        ( $tt:ty ) => {
+            impl ToInt4BE for $tt {
+                fn convert_to_i4be(ii: $tt) -> Result<$tt, $tt> {
+                    if ii < i32::MIN as $tt { return Err(ii); }
+                    if ii > i32::MAX as $tt { return Err(ii); }
+
+                    let mut buf = [0; size_of::<$tt>()];
+                    BigEndian::write_i32(&mut buf, ii as i32);
+                    Ok(<$tt>::from_le_bytes(buf))
+                }
+            }
+        }
+    }
+
+    impl_i2be!(f64);
+    impl_i4be!(f64);
+
+    impl_i2be!(f32);
+    impl_i4be!(f32);
+
+    impl_i2be!(i64);
+    impl_i4be!(i64);
+    impl_i2be!(u64);
+    impl_i4be!(u64);
+
+    impl_i2be!(i32);
+    impl_i4be!(i32);
+    impl_i2be!(u32);
+    impl_i4be!(u32);
+
+    impl_i2be!(i16);
+    impl_i2be!(u16);
+
+    // Does not fit
+    //impl_i4be!(i16);
+    //impl_i4be!(u16);
+    //
+    //impl_i2be!(i8);
+    //impl_i4be!(i8);
+    //impl_i2be!(u8);
+    //impl_i4be!(u8);
+
+}