pyo3 variant working

Cargo.toml

@@ -1,9 +1,16 @@
 [package]
-name = "rs-klamath"
+name = "klamath_rs_ext"
 version = "0.1.0"
 authors = ["jan <jan@mpxd.net>"]
 edition = "2021"
 
+
+[lib]
+name = "klamath_rs_ext"
+crate-type = ["cdylib", "rlib"]
+
+
 [dependencies]
 byteorder = "^1"
-#nom = "^7"
+pyo3 = "^0"
+numpy = "^0"

klamath_rs_ext/__init__.py

@@ -0,0 +1,6 @@
+from .basic import pack_int2 as pack_int2
+from .basic import pack_int4 as pack_int4
+
+
+__version__ = 0.1
+

klamath_rs_ext/basic.py

@@ -0,0 +1,56 @@
+from collections.abc import Sequence
+
+import numpy
+from numpy.typing import NDArray
+
+from .klamath_rs_ext import arr_to_int2, arr_to_int4
+
+
+def pack_int2(data: NDArray[numpy.integer] | Sequence[int] | int) -> bytes:
+    arr = numpy.asarray(data)
+
+    if arr.dtype in (
+            numpy.float64, numpy.float32,
+            numpy.int64, numpy.uint64,
+            numpy.int32, numpy.uint32,
+            numpy.int16, numpy.uint16,
+            ):
+        arr = numpy.require(arr, requirements=('C_CONTIGUOUS', 'ALIGNED', 'WRITEABLE', 'OWNDATA'))
+        if arr is data:
+            arr = numpy.array(arr, copy=True)
+        arr_to_int2(arr)
+        i2arr = arr.view('>i2')[::arr.itemsize // 2]
+        return i2arr.tobytes()
+
+    if arr.dtype == numpy.dtype('>i2'):
+        return arr.tobytes()
+
+    if (arr > 32767).any() or (arr < -32768).any():
+        raise Exception(f'int2 data out of range: {arr}')
+
+    return arr.astype('>i2').tobytes()
+
+
+def pack_int4(data: NDArray[numpy.integer] | Sequence[int] | int) -> bytes:
+    arr = numpy.asarray(data)
+
+    if arr.dtype in (
+            numpy.float64, numpy.float32,
+            numpy.int64, numpy.uint64,
+            numpy.int32, numpy.uint32,
+            ):
+        arr = numpy.require(arr, requirements=('C_CONTIGUOUS', 'ALIGNED', 'WRITEABLE', 'OWNDATA'))
+        if arr is data:
+            arr = numpy.array(arr, copy=True)
+        arr_to_int4(arr)
+        i4arr = arr.view('>i4')[::arr.itemsize // 4]
+        return i4arr.tobytes()
+
+    if arr.dtype == numpy.dtype('>i4'):
+        return arr.tobytes()
+
+    if (arr > 2147483647).any() or (arr < -2147483648).any():
+        raise Exception(f'int4 data out of range: {arr}')
+
+    return arr.astype('>i4').tobytes()
+

pyproject.toml

@@ -0,0 +1,7 @@
+[build-system]
+requires = ["maturin>1.0,<2.0"]
+build-backend = "maturin"
+
+
+[tool.maturin]
+features = ["pyo3/extension-module"]

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);
+
+}