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

//! PyGroup — navigable HDF5 group with read and write support.

use std::collections::HashMap;
use std::sync::{Arc, Mutex, OnceLock};

use pyo3::exceptions::{PyIOError, PyKeyError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::PyList;

use crate::attrs::PyAttrs;
use crate::{DatasetSpec, OwnedAttrValue, apply_dataset_spec, extract_numpy_data, node};

/// Shared state for a group being written.
pub(crate) struct WriteGroupState {
    pub name: String,
    pub datasets: Vec<DatasetSpec>,
    pub attrs: Arc<Mutex<Vec<(String, OwnedAttrValue)>>>,
}

/// An HDF5 group.
///
/// In read mode it behaves like an h5py group: `grp['name']`,
/// `grp['sub/path']` and `grp['/absolute/path']`, `keys()`, `values()`,
/// `items()`, iteration, `len()`, `in`, `get()`, `name` and `attrs`.
/// In write mode, supports `create_dataset` and attribute setting.
#[pyclass(name = "Group")]
pub struct PyGroup {
    inner: GroupInner,
}

enum GroupInner {
    Read(ReadGroup),
    Write(Arc<Mutex<WriteGroupState>>),
}

impl PyGroup {
    pub(crate) fn from_read(file: Arc<clawhdf5_rs::File>, path: String, addr: u64) -> Self {
        Self {
            inner: GroupInner::Read(ReadGroup::new(file, path, addr)),
        }
    }

    pub(crate) fn from_write(state: Arc<Mutex<WriteGroupState>>) -> Self {
        Self {
            inner: GroupInner::Write(state),
        }
    }

    fn read_group(&self, what: &str) -> PyResult<&ReadGroup> {
        match &self.inner {
            GroupInner::Read(g) => Ok(g),
            GroupInner::Write(_) => Err(PyIOError::new_err(format!(
                "cannot {what} a group opened for writing"
            ))),
        }
    }
}

/// A group in a file opened for reading (a file is its root group, as in
/// h5py). It keeps its own address and, once listed, its links, so looking
/// up a child neither resolves the path from the root nor scans the group's
/// links again: visiting every member of a large group is linear, not
/// quadratic.
pub(crate) struct ReadGroup {
    pub file: Arc<clawhdf5_rs::File>,
    pub path: String,
    pub addr: u64,
    /// Link name -> object address (soft links resolved), filled on first use.
    links: OnceLock<HashMap<String, u64>>,
    /// Names of the datasets and subgroups, sorted (h5py's order).
    members: OnceLock<Vec<String>>,
}

impl ReadGroup {
    pub(crate) fn new(file: Arc<clawhdf5_rs::File>, path: String, addr: u64) -> Self {
        Self {
            file,
            path,
            addr,
            links: OnceLock::new(),
            members: OnceLock::new(),
        }
    }

    fn links(&self) -> PyResult<&HashMap<String, u64>> {
        if let Some(links) = self.links.get() {
            return Ok(links);
        }
        let entries = crate::no_panic(|| {
            clawhdf5_format::group_v2::resolve_group_children(
                self.file.as_bytes(),
                self.file.superblock(),
                self.addr,
            )
            .map_err(|e| PyValueError::new_err(format!("{}: {e}", node::name(&self.path))))
        })?;
        let map = entries
            .into_iter()
            .map(|e| (e.name, e.object_header_address))
            .collect();
        Ok(self.links.get_or_init(|| map))
    }

    /// The path and address of `key` (a name, a relative or an absolute path).
    fn locate(&self, key: &str) -> PyResult<(String, u64)> {
        let path = node::join(&self.path, key);
        let rel = if self.path.is_empty() {
            Some(path.as_str())
        } else if path == self.path {
            Some("")
        } else {
            path.strip_prefix(self.path.as_str())
                .and_then(|r| r.strip_prefix('/'))
        };
        let addr = match rel {
            // A direct child: the link table, when it has the name.
            Some(name) if !name.is_empty() && !name.contains('/') => {
                match self.links()?.get(name) {
                    Some(&a) => a,
                    None => node::resolve_from(&self.file, self.addr, name, &path)?,
                }
            }
            Some(rel) => node::resolve_from(&self.file, self.addr, rel, &path)?,
            None => node::address(&self.file, &path)?,
        };
        Ok((path, addr))
    }

    /// `group[key]`.
    pub(crate) fn get_item(&self, py: Python<'_>, key: &str) -> PyResult<Py<PyAny>> {
        let (path, addr) = self.locate(key)?;
        node::open(py, &self.file, path, addr)
    }

    /// `group.get(key, default)`.
    pub(crate) fn get(
        &self,
        py: Python<'_>,
        key: &str,
        default: Option<Py<PyAny>>,
    ) -> PyResult<Py<PyAny>> {
        match self.get_item(py, key) {
            Err(e) if e.is_instance_of::<PyKeyError>(py) => {
                Ok(default.unwrap_or_else(|| py.None()))
            }
            other => other,
        }
    }

    /// Names of the group's datasets and subgroups, sorted (h5py's order).
    pub(crate) fn member_names(&self) -> PyResult<&[String]> {
        if let Some(m) = self.members.get() {
            return Ok(m);
        }
        let mut names = Vec::new();
        for (name, &addr) in self.links()? {
            let hdr = node::header_at(&self.file, addr, &node::join(&self.path, name))?;
            if matches!(
                node::kind(&hdr),
                Some(node::Kind::Dataset | node::Kind::Group)
            ) {
                names.push(name.clone());
            }
        }
        names.sort_by(|a, b| a.as_bytes().cmp(b.as_bytes()));
        Ok(self.members.get_or_init(|| names))
    }

    pub(crate) fn contains(&self, key: &str) -> bool {
        self.locate(key)
            .and_then(|(path, addr)| node::header_at(&self.file, addr, &path))
            .ok()
            .and_then(|h| node::kind(&h))
            .is_some_and(|k| k != node::Kind::Datatype)
    }

    pub(crate) fn values(&self, py: Python<'_>) -> PyResult<Vec<Py<PyAny>>> {
        self.member_names()?
            .iter()
            .map(|n| self.get_item(py, n))
            .collect()
    }

    pub(crate) fn items(&self, py: Python<'_>) -> PyResult<Vec<(String, Py<PyAny>)>> {
        self.member_names()?
            .iter()
            .map(|n| Ok((n.clone(), self.get_item(py, n)?)))
            .collect()
    }

    pub(crate) fn attrs(&self) -> PyResult<PyAttrs> {
        PyAttrs::read(Arc::clone(&self.file), self.addr, &self.path)
    }
}

#[pymethods]
impl PyGroup {
    /// Get a child object (dataset or subgroup) by name or path.
    fn __getitem__(&self, py: Python<'_>, key: &str) -> PyResult<Py<PyAny>> {
        self.read_group("read children from")?.get_item(py, key)
    }

    /// `group.get(key, default=None)`.
    #[pyo3(signature = (key, default=None))]
    fn get(&self, py: Python<'_>, key: &str, default: Option<Py<PyAny>>) -> PyResult<Py<PyAny>> {
        self.read_group("read children from")?.get(py, key, default)
    }

    /// List the names of all children (datasets and subgroups).
    fn keys(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        match &self.inner {
            GroupInner::Read(g) => {
                let list = PyList::new(py, g.member_names()?)?;
                Ok(list.into_any().unbind())
            }
            GroupInner::Write(state) => {
                let guard = state.lock().unwrap();
                let names: Vec<&str> = guard.datasets.iter().map(|d| d.name.as_str()).collect();
                let list = PyList::new(py, &names)?;
                Ok(list.into_any().unbind())
            }
        }
    }

    fn values(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        let g = self.read_group("read children from")?;
        Ok(PyList::new(py, g.values(py)?)?.into_any().unbind())
    }

    fn items(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        let g = self.read_group("read children from")?;
        Ok(PyList::new(py, g.items(py)?)?.into_any().unbind())
    }

    fn __iter__(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        self.keys(py)?.call_method0(py, "__iter__")
    }

    fn __len__(&self) -> PyResult<usize> {
        match &self.inner {
            GroupInner::Read(g) => Ok(g.member_names()?.len()),
            GroupInner::Write(state) => Ok(state.lock().unwrap().datasets.len()),
        }
    }

    /// The group's full name, e.g. `/sensors`.
    #[getter]
    fn name(&self) -> String {
        match &self.inner {
            GroupInner::Read(g) => node::name(&g.path),
            GroupInner::Write(state) => node::name(&state.lock().unwrap().name),
        }
    }

    /// Create a dataset inside this group (write mode only).
    ///
    /// Parameters:
    ///   name: dataset name
    ///   data: numpy array
    ///   chunks: optional chunk dimensions
    ///   compression: optional, only 'gzip' supported
    ///   compression_opts: gzip level (1-9)
    #[pyo3(signature = (name, *, data, chunks=None, compression=None, compression_opts=None))]
    fn create_dataset(
        &self,
        py: Python<'_>,
        name: &str,
        data: &Bound<'_, PyAny>,
        chunks: Option<Vec<u64>>,
        compression: Option<&str>,
        compression_opts: Option<u32>,
    ) -> PyResult<()> {
        match &self.inner {
            GroupInner::Write(state) => {
                let (dataset_data, shape) = extract_numpy_data(py, data)?;
                let deflate_level = match compression {
                    Some("gzip") => Some(compression_opts.unwrap_or(4)),
                    Some(other) => {
                        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
                            "unsupported compression: {other}; only 'gzip' is supported"
                        )));
                    }
                    None => None,
                };
                let spec = DatasetSpec {
                    name: name.to_string(),
                    data: dataset_data,
                    shape,
                    chunks,
                    deflate_level,
                    attrs: vec![],
                };
                state.lock().unwrap().datasets.push(spec);
                Ok(())
            }
            GroupInner::Read { .. } => Err(PyIOError::new_err(
                "cannot create datasets on a read-only group",
            )),
        }
    }

    /// Attribute access.
    #[getter]
    fn attrs(&self) -> PyResult<PyAttrs> {
        match &self.inner {
            GroupInner::Read(g) => g.attrs(),
            GroupInner::Write(state) => {
                let store = Arc::clone(&state.lock().unwrap().attrs);
                Ok(PyAttrs::from_write(store))
            }
        }
    }

    fn __repr__(&self) -> String {
        match &self.inner {
            GroupInner::Read(g) => {
                let n = g.member_names().map_or(0, |m| m.len());
                format!("<HDF5 group \"{}\" ({n} members)>", node::name(&g.path))
            }
            GroupInner::Write(state) => {
                let name = &state.lock().unwrap().name;
                format!("<HDF5 Group \"{name}\" (write)>")
            }
        }
    }

    fn __contains__(&self, key: &str) -> PyResult<bool> {
        match &self.inner {
            GroupInner::Read(g) => Ok(g.contains(key)),
            GroupInner::Write(state) => {
                let guard = state.lock().unwrap();
                Ok(guard.datasets.iter().any(|d| d.name == key))
            }
        }
    }
}

/// Finalize a write group into the file builder.
pub(crate) fn finalize_write_group(
    builder: &mut clawhdf5_rs::FileBuilder,
    state: &WriteGroupState,
) {
    let mut gb = builder.create_group(&state.name);
    for spec in &state.datasets {
        let db = gb.create_dataset(&spec.name);
        apply_dataset_spec(db, spec);
    }
    let attrs_guard = state.attrs.lock().unwrap();
    for (name, val) in attrs_guard.iter() {
        gb.set_attr(name, val.clone().into());
    }
    let finished = gb.finish();
    builder.add_group(finished);
}

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

    #[test]
    fn member_names_are_sorted() {
        let mut b = clawhdf5_rs::FileBuilder::new();
        b.create_dataset("zeta").with_f64_data(&[1.0]);
        b.create_dataset("alpha").with_f64_data(&[1.0]);
        let mut g = b.create_group("mid");
        g.create_dataset("x").with_f64_data(&[1.0]);
        let finished = g.finish();
        b.add_group(finished);
        let bytes = b.finish().unwrap();
        let file = Arc::new(clawhdf5_rs::File::from_bytes(bytes).unwrap());
        let root = file.superblock().root_group_address;
        let top = ReadGroup::new(Arc::clone(&file), String::new(), root);
        assert_eq!(top.member_names().unwrap(), ["alpha", "mid", "zeta"]);
        let (path, addr) = top.locate("mid").unwrap();
        assert_eq!(path, "mid");
        let mid = ReadGroup::new(Arc::clone(&file), path, addr);
        assert_eq!(mid.member_names().unwrap(), ["x"]);
        assert!(top.contains("mid/x"));
        assert!(mid.contains("/alpha"));
        assert!(mid.contains("x") && mid.contains("./x"));
        assert!(!top.contains("nope"));
        assert!(!mid.contains("alpha"));
    }

    #[test]
    fn finalize_group() {
        let state = WriteGroupState {
            name: "mygroup".into(),
            datasets: vec![DatasetSpec {
                name: "vals".into(),
                data: crate::DatasetData::F64(vec![1.0, 2.0]),
                shape: vec![2],
                chunks: None,
                deflate_level: None,
                attrs: vec![],
            }],
            attrs: Arc::new(Mutex::new(vec![("version".into(), OwnedAttrValue::I64(1))])),
        };
        let mut builder = clawhdf5_rs::FileBuilder::new();
        // Need a root dataset for a valid file
        builder.create_dataset("root_ds").with_f64_data(&[0.0]);
        finalize_write_group(&mut builder, &state);
        let bytes = builder.finish().unwrap();
        let file = clawhdf5_rs::File::from_bytes(bytes).unwrap();
        let ds = file.dataset("mygroup/vals").unwrap();
        assert_eq!(ds.read_f64().unwrap(), vec![1.0, 2.0]);
    }
}