From ba07d253d266b0eb19b64e2663975b772b7198c0 Mon Sep 17 00:00:00 2001 From: jan Date: Sat, 21 Dec 2024 09:50:19 -0800 Subject: [PATCH] pyo3 variant working --- Cargo.toml | 11 ++- klamath_rs_ext/__init__.py | 6 ++ klamath_rs_ext/basic.py | 56 +++++++++++++ klamath_rs_ext/py.typed | 0 pyproject.toml | 7 ++ src/lib.rs | 156 +++++++++++++++++++++++++++++++++++++ 6 files changed, 234 insertions(+), 2 deletions(-) create mode 100644 klamath_rs_ext/__init__.py create mode 100644 klamath_rs_ext/basic.py create mode 100644 klamath_rs_ext/py.typed create mode 100644 pyproject.toml diff --git a/Cargo.toml b/Cargo.toml index 01429e9..41cbb68 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,9 +1,16 @@ [package] -name = "rs-klamath" +name = "klamath_rs_ext" version = "0.1.0" authors = ["jan "] edition = "2021" + +[lib] +name = "klamath_rs_ext" +crate-type = ["cdylib", "rlib"] + + [dependencies] byteorder = "^1" -#nom = "^7" +pyo3 = "^0" +numpy = "^0" diff --git a/klamath_rs_ext/__init__.py b/klamath_rs_ext/__init__.py new file mode 100644 index 0000000..fa02fb7 --- /dev/null +++ b/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 + diff --git a/klamath_rs_ext/basic.py b/klamath_rs_ext/basic.py new file mode 100644 index 0000000..7237b1c --- /dev/null +++ b/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() + diff --git a/klamath_rs_ext/py.typed b/klamath_rs_ext/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..9704ea5 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,7 @@ +[build-system] +requires = ["maturin>1.0,<2.0"] +build-backend = "maturin" + + +[tool.maturin] +features = ["pyo3/extension-module"] diff --git a/src/lib.rs b/src/lib.rs index 2a34034..b64ab5d 100644 --- a/src/lib.rs +++ b/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::>()?; + 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::>()?; + 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 where Self: Sized; + } + + pub trait ToInt4BE { + fn convert_to_i4be(ii: Self) -> Result 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); + +}