klamath-rs/src/lib.rs

extern crate byteorder;

pub mod basic;
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);

}