clawhdf5/crates/clawhdf5-py/src/lib.rs

//! Python bindings for clawhdf5 — a pure-Rust HDF5 library.
//!
//! Provides a Pythonic API mirroring h5py:
//!
//! ```python
//! import clawhdf5
//!
//! with clawhdf5.File('data.h5', 'r') as f:
//!     data = f['dataset_name'][:]
//!
//! with clawhdf5.File('http://host/data.h5') as f:   # range requests
//!     block = f['dataset_name'][10:20]
//!
//! with clawhdf5.File('data.h5', 'r+') as f:          # in-place edits
//!     f['dataset_name'][0] = 1.5
//!     f.attrs['note'] = 'edited'
//! ```

mod attrs;
mod convert;
mod dataset;
mod edit;
mod file;
mod group;
mod handle;
mod node;
mod select;

use pyo3::prelude::*;

pub(crate) use attrs::PyAttrs;
pub(crate) use dataset::PyDataset;
pub(crate) use file::PyFile;
pub(crate) use group::PyGroup;

pyo3::create_exception!(
    clawhdf5,
    InternalError,
    pyo3::exceptions::PyRuntimeError,
    "A bug in clawhdf5 met while reading or writing a file (a Rust panic, \
     caught). Derived from RuntimeError, so `except Exception` handles it."
);

/// The text of a caught panic.
pub(crate) fn panic_text(payload: &(dyn std::any::Any + Send)) -> String {
    payload
        .downcast_ref::<&str>()
        .map(|s| (*s).to_string())
        .or_else(|| payload.downcast_ref::<String>().cloned())
        .unwrap_or_else(|| "unknown panic".to_string())
}

/// Run `f`, turning a panic in the library into [`InternalError`] instead of
/// PyO3's `PanicException` (a `BaseException`, which `except Exception`
/// does not catch). Wraps every call into the library.
pub(crate) fn no_panic<T>(f: impl FnOnce() -> PyResult<T>) -> PyResult<T> {
    std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)).unwrap_or_else(|p| {
        Err(InternalError::new_err(format!(
            "clawhdf5 internal error (please report it): {}",
            panic_text(&*p)
        )))
    })
}

/// A test hook: panics inside [`no_panic`], so the tests can check that a
/// library panic reaches Python as an ordinary exception.
#[pyfunction]
fn _panic_for_test() -> PyResult<()> {
    no_panic(|| panic!("deliberate panic for the test suite"))
}

/// Convert a `clawhdf5_rs::Error` into a `PyErr`.
///
/// Maps different error variants to more specific Python exception types:
/// - I/O errors, and failed reads of a remote file -> `PyIOError`/`PyOSError`
/// - Format/parsing errors -> `PyValueError`
/// - Missing dataset/path errors -> `PyKeyError`
/// - Invalid arguments -> `PyValueError`
/// - Unsupported operations -> `PyNotImplementedError`
/// - Other errors (a locked file, ...) -> `PyOSError`
pub(crate) fn to_py_err(e: clawhdf5_rs::Error) -> PyErr {
    use clawhdf5_rs::Error;
    match &e {
        Error::Io(_) => PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()),
        // A failed read of the storage: a network error on a remote file.
        Error::Format(clawhdf5_format::error::FormatError::Storage(_)) => {
            PyErr::new::<pyo3::exceptions::PyOSError, _>(e.to_string())
        }
        Error::Format(_) => PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()),
        Error::NotADataset(_) | Error::MissingMessage(_) => {
            PyErr::new::<pyo3::exceptions::PyKeyError, _>(e.to_string())
        }
        Error::AlignmentError(_)
        | Error::ZeroCopyNotContiguous
        | Error::ZeroCopyNonNativeEndian
        | Error::ZeroCopyTypeMismatch { .. }
        | Error::ZeroCopyUnaligned { .. }
        | Error::InvalidArgument(_) => {
            PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string())
        }
        Error::Unsupported(_) => {
            PyErr::new::<pyo3::exceptions::PyNotImplementedError, _>(e.to_string())
        }
        _ => PyErr::new::<pyo3::exceptions::PyOSError, _>(e.to_string()),
    }
}

/// The value of a dataset or attribute with a null dataspace: a type but no
/// data. Mirrors `h5py.Empty`.
#[pyclass(name = "Empty", frozen)]
pub struct PyEmpty {
    dtype: Py<PyAny>,
}

impl PyEmpty {
    pub(crate) fn new(dtype: Py<PyAny>) -> Self {
        Self { dtype }
    }
}

#[pymethods]
impl PyEmpty {
    #[new]
    fn py_new(py: Python<'_>, dtype: &Bound<'_, PyAny>) -> PyResult<Self> {
        let dtype = py.import("numpy")?.getattr("dtype")?.call1((dtype,))?;
        Ok(Self::new(dtype.unbind()))
    }

    #[getter]
    fn dtype(&self, py: Python<'_>) -> Py<PyAny> {
        self.dtype.clone_ref(py)
    }

    #[getter]
    fn shape(&self, py: Python<'_>) -> Py<PyAny> {
        py.None()
    }

    #[getter]
    fn size(&self, py: Python<'_>) -> Py<PyAny> {
        py.None()
    }

    fn __eq__(&self, py: Python<'_>, other: &Bound<'_, PyAny>) -> PyResult<bool> {
        match other.cast::<PyEmpty>() {
            Ok(o) => self.dtype.bind(py).eq(o.get().dtype.bind(py)),
            Err(_) => Ok(false),
        }
    }

    fn __repr__(&self, py: Python<'_>) -> PyResult<String> {
        Ok(format!("Empty(dtype={})", self.dtype.bind(py).repr()?))
    }
}

/// The data payload for a dataset being written.
#[derive(Clone)]
pub(crate) enum DatasetData {
    F64(Vec<f64>),
    F32(Vec<f32>),
    I64(Vec<i64>),
    I32(Vec<i32>),
    U8(Vec<u8>),
}

/// Specification for a dataset to be written.
#[derive(Clone)]
pub(crate) struct DatasetSpec {
    pub name: String,
    pub data: DatasetData,
    pub shape: Vec<u64>,
    pub chunks: Option<Vec<u64>>,
    pub deflate_level: Option<u32>,
    pub attrs: Vec<(String, OwnedAttrValue)>,
}

/// Owned attribute value (used during write accumulation).
#[derive(Clone)]
pub(crate) enum OwnedAttrValue {
    F64(f64),
    I64(i64),
    Str(String),
    F64Array(Vec<f64>),
    I64Array(Vec<i64>),
}

impl From<OwnedAttrValue> for clawhdf5_rs::AttrValue {
    fn from(v: OwnedAttrValue) -> Self {
        match v {
            OwnedAttrValue::F64(x) => clawhdf5_rs::AttrValue::F64(x),
            OwnedAttrValue::I64(x) => clawhdf5_rs::AttrValue::I64(x),
            OwnedAttrValue::Str(s) => clawhdf5_rs::AttrValue::String(s),
            OwnedAttrValue::F64Array(a) => clawhdf5_rs::AttrValue::F64Array(a),
            OwnedAttrValue::I64Array(a) => clawhdf5_rs::AttrValue::I64Array(a),
        }
    }
}

/// Extract a Python value into an `OwnedAttrValue`.
pub(crate) fn py_to_attr_value(val: &Bound<'_, PyAny>) -> PyResult<OwnedAttrValue> {
    // Try int first (before float, since bool is int subclass in Python)
    if let Ok(v) = val.extract::<i64>() {
        return Ok(OwnedAttrValue::I64(v));
    }
    if let Ok(v) = val.extract::<f64>() {
        return Ok(OwnedAttrValue::F64(v));
    }
    if let Ok(v) = val.extract::<String>() {
        return Ok(OwnedAttrValue::Str(v));
    }
    // Try list of floats, list of ints
    if let Ok(v) = val.extract::<Vec<f64>>() {
        return Ok(OwnedAttrValue::F64Array(v));
    }
    if let Ok(v) = val.extract::<Vec<i64>>() {
        return Ok(OwnedAttrValue::I64Array(v));
    }
    Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
        "unsupported attribute type; expected int, float, str, or list of int/float",
    ))
}

/// Convert an `AttrValue` (from the Rust lib) to a Python object.
pub(crate) fn attr_value_to_py(py: Python<'_>, val: &clawhdf5_rs::AttrValue) -> Py<PyAny> {
    match val {
        clawhdf5_rs::AttrValue::F64(v) => v.into_pyobject(py).unwrap().into_any().unbind(),
        clawhdf5_rs::AttrValue::I64(v) => v.into_pyobject(py).unwrap().into_any().unbind(),
        clawhdf5_rs::AttrValue::U64(v) => v.into_pyobject(py).unwrap().into_any().unbind(),
        clawhdf5_rs::AttrValue::String(s) => s.into_pyobject(py).unwrap().into_any().unbind(),
        clawhdf5_rs::AttrValue::F64Array(a) => {
            let list = pyo3::types::PyList::new(py, a).unwrap();
            list.into_any().unbind()
        }
        clawhdf5_rs::AttrValue::I64Array(a) => {
            let list = pyo3::types::PyList::new(py, a).unwrap();
            list.into_any().unbind()
        }
        clawhdf5_rs::AttrValue::U64Array(a) => {
            let list = pyo3::types::PyList::new(py, a).unwrap();
            list.into_any().unbind()
        }
        clawhdf5_rs::AttrValue::StringArray(a) => {
            let list = pyo3::types::PyList::new(py, a).unwrap();
            list.into_any().unbind()
        }
        // No Python-side decoding for this datatype: hand back everything
        // needed to interpret it rather than dropping the attribute.
        clawhdf5_rs::AttrValue::Raw {
            datatype,
            shape,
            data,
        } => {
            let dict = pyo3::types::PyDict::new(py);
            dict.set_item("dtype", format!("{datatype:?}")).unwrap();
            dict.set_item("shape", shape).unwrap();
            dict.set_item("data", pyo3::types::PyBytes::new(py, data))
                .unwrap();
            dict.into_any().unbind()
        }
    }
}

/// Apply a `DatasetSpec` to a `DatasetBuilder`.
pub(crate) fn apply_dataset_spec(
    db: &mut clawhdf5_format::type_builders::DatasetBuilder,
    spec: &DatasetSpec,
) {
    match &spec.data {
        DatasetData::F64(v) => {
            db.with_f64_data(v);
        }
        DatasetData::F32(v) => {
            db.with_f32_data(v);
        }
        DatasetData::I64(v) => {
            db.with_i64_data(v);
        }
        DatasetData::I32(v) => {
            db.with_i32_data(v);
        }
        DatasetData::U8(v) => {
            db.with_u8_data(v);
        }
    }
    if !spec.shape.is_empty() {
        db.with_shape(&spec.shape);
    }
    if let Some(chunks) = &spec.chunks {
        db.with_chunks(chunks);
    }
    if let Some(level) = spec.deflate_level {
        db.with_deflate(level);
    }
    for (name, val) in &spec.attrs {
        db.set_attr(name, val.clone().into());
    }
}

/// Extract numpy array data from a Python object.
pub(crate) fn extract_numpy_data(
    py: Python<'_>,
    data: &Bound<'_, PyAny>,
) -> PyResult<(DatasetData, Vec<u64>)> {
    let np = py.import("numpy")?;
    let arr = np.call_method1("ascontiguousarray", (data,))?;
    let dtype_str: String = arr.getattr("dtype")?.str()?.extract()?;
    let shape: Vec<usize> = arr.getattr("shape")?.extract()?;
    let shape_u64: Vec<u64> = shape.iter().map(|&s| s as u64).collect();
    let flat = arr.call_method0("ravel")?;

    let dataset_data = match dtype_str.as_str() {
        "float64" => DatasetData::F64(flat.extract::<Vec<f64>>()?),
        "float32" => DatasetData::F32(flat.extract::<Vec<f32>>()?),
        "int64" => DatasetData::I64(flat.extract::<Vec<i64>>()?),
        "int32" => DatasetData::I32(flat.extract::<Vec<i32>>()?),
        "uint8" => DatasetData::U8(flat.extract::<Vec<u8>>()?),
        _ => {
            return Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(format!(
                "unsupported numpy dtype: {dtype_str}; expected float64, float32, int64, int32, or uint8"
            )));
        }
    };
    Ok((dataset_data, shape_u64))
}

/// The clawhdf5 Python module.
#[pymodule]
fn clawhdf5(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
    m.add_class::<PyFile>()?;
    m.add_class::<PyDataset>()?;
    m.add_class::<PyGroup>()?;
    m.add_class::<PyAttrs>()?;
    m.add_class::<PyEmpty>()?;
    m.add("InternalError", m.py().get_type::<InternalError>())?;
    m.add_function(wrap_pyfunction!(_panic_for_test, m)?)?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn owned_attr_value_roundtrip() {
        let val = OwnedAttrValue::F64(2.5);
        let attr: clawhdf5_rs::AttrValue = val.into();
        assert!(matches!(attr, clawhdf5_rs::AttrValue::F64(v) if (v - 2.5).abs() < 1e-10));
    }

    #[test]
    fn owned_attr_value_i64() {
        let val = OwnedAttrValue::I64(42);
        let attr: clawhdf5_rs::AttrValue = val.into();
        assert!(matches!(attr, clawhdf5_rs::AttrValue::I64(42)));
    }

    #[test]
    fn owned_attr_value_str() {
        let val = OwnedAttrValue::Str("hello".into());
        let attr: clawhdf5_rs::AttrValue = val.into();
        assert!(matches!(attr, clawhdf5_rs::AttrValue::String(ref s) if s == "hello"));
    }

    #[test]
    fn owned_attr_value_f64_array() {
        let val = OwnedAttrValue::F64Array(vec![1.0, 2.0]);
        let attr: clawhdf5_rs::AttrValue = val.into();
        assert!(matches!(attr, clawhdf5_rs::AttrValue::F64Array(ref v) if v == &[1.0, 2.0]));
    }

    #[test]
    fn owned_attr_value_i64_array() {
        let val = OwnedAttrValue::I64Array(vec![1, 2, 3]);
        let attr: clawhdf5_rs::AttrValue = val.into();
        assert!(matches!(attr, clawhdf5_rs::AttrValue::I64Array(ref v) if v == &[1, 2, 3]));
    }

    #[test]
    fn dataset_spec_apply() {
        let spec = DatasetSpec {
            name: "test".into(),
            data: DatasetData::F64(vec![1.0, 2.0, 3.0]),
            shape: vec![3],
            chunks: None,
            deflate_level: None,
            attrs: vec![],
        };
        let mut builder = clawhdf5_rs::FileBuilder::new();
        let db = builder.create_dataset(&spec.name);
        apply_dataset_spec(db, &spec);
        let bytes = builder.finish().unwrap();
        let file = clawhdf5_rs::File::from_bytes(bytes).unwrap();
        let ds = file.dataset("test").unwrap();
        assert_eq!(ds.read_f64().unwrap(), vec![1.0, 2.0, 3.0]);
    }

    #[test]
    fn dataset_spec_with_chunks_and_deflate() {
        let spec = DatasetSpec {
            name: "compressed".into(),
            data: DatasetData::I32(vec![10, 20, 30, 40, 50]),
            shape: vec![5],
            chunks: Some(vec![5]),
            deflate_level: Some(4),
            attrs: vec![("unit".into(), OwnedAttrValue::Str("m".into()))],
        };
        let mut builder = clawhdf5_rs::FileBuilder::new();
        let db = builder.create_dataset(&spec.name);
        apply_dataset_spec(db, &spec);
        let bytes = builder.finish().unwrap();
        let file = clawhdf5_rs::File::from_bytes(bytes).unwrap();
        let ds = file.dataset("compressed").unwrap();
        assert_eq!(ds.read_i32().unwrap(), vec![10, 20, 30, 40, 50]);
    }
}