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

//! Resolving paths to objects in a file opened for reading.

use std::sync::Arc;

use clawhdf5_format::attribute::AttributeMessage;
use clawhdf5_format::dataspace::{Dataspace, DataspaceType};
use clawhdf5_format::message_type::MessageType;
use clawhdf5_format::object_header::ObjectHeader;
use pyo3::exceptions::{PyKeyError, PyTypeError, PyValueError};
use pyo3::prelude::*;

use crate::dataset::PyDataset;
use crate::group::PyGroup;

/// Join `key` onto the group path `base` the way h5py does: an absolute key
/// starts from the root, a relative one from `base`. Paths are kept without
/// a leading `/`; the root is `""`.
pub(crate) fn join(base: &str, key: &str) -> String {
    let parts = if key.starts_with('/') {
        key.split('/').collect::<Vec<_>>()
    } else {
        base.split('/').chain(key.split('/')).collect()
    };
    parts
        .into_iter()
        .filter(|p| !p.is_empty() && *p != ".")
        .collect::<Vec<_>>()
        .join("/")
}

/// The HDF5 name (`/a/b`) of a path.
pub(crate) fn name(path: &str) -> String {
    format!("/{path}")
}

/// The address of the object at `path`, resolved from the root group.
pub(crate) fn address(file: &clawhdf5_rs::File, path: &str) -> PyResult<u64> {
    resolve_from(file, file.superblock().root_group_address, path, path)
}

/// The address of `rel` resolved from the group at `group` (`full` is the
/// resulting path, for the error message).
pub(crate) fn resolve_from(
    file: &clawhdf5_rs::File,
    group: u64,
    rel: &str,
    full: &str,
) -> PyResult<u64> {
    if rel.is_empty() {
        return Ok(group);
    }
    crate::no_panic(|| {
        clawhdf5_format::group_v2::resolve_path_from(file.as_bytes(), file.superblock(), group, rel)
            .map_err(|e| {
                PyKeyError::new_err(format!(
                    "Unable to open object (object '{}' doesn't exist): {e}",
                    name(full)
                ))
            })
    })
}

/// The object header at `addr` (the object at `path`).
pub(crate) fn header_at(file: &clawhdf5_rs::File, addr: u64, path: &str) -> PyResult<ObjectHeader> {
    crate::no_panic(|| {
        let sb = file.superblock();
        let at = usize::try_from(addr)
            .map_err(|_| PyValueError::new_err(format!("{}: address out of range", name(path))))?;
        ObjectHeader::parse(file.as_bytes(), at, sb.offset_size, sb.length_size)
            .map_err(|e| PyValueError::new_err(format!("{}: {e}", name(path))))
    })
}

/// What an object header describes.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Kind {
    Dataset,
    Group,
    Datatype,
}

pub(crate) fn kind(hdr: &ObjectHeader) -> Option<Kind> {
    let has = |t: MessageType| hdr.messages.iter().any(|m| m.msg_type == t);
    if has(MessageType::DataLayout) {
        Some(Kind::Dataset)
    } else if has(MessageType::LinkInfo)
        || has(MessageType::Link)
        || has(MessageType::SymbolTable)
        || has(MessageType::GroupInfo)
    {
        Some(Kind::Group)
    } else if has(MessageType::Datatype) {
        Some(Kind::Datatype)
    } else {
        None
    }
}

/// Open the object at `addr` (whose path is `path`) as a `Dataset` or
/// `Group`. Both keep the address, so later reads resolve nothing.
pub(crate) fn open(
    py: Python<'_>,
    file: &Arc<clawhdf5_rs::File>,
    path: String,
    addr: u64,
) -> PyResult<Py<PyAny>> {
    let hdr = header_at(file, addr, &path)?;
    match kind(&hdr) {
        Some(Kind::Dataset) => Ok(PyDataset::open(py, Arc::clone(file), path, addr, &hdr)?
            .into_pyobject(py)?
            .into_any()
            .unbind()),
        Some(Kind::Group) => Ok(PyGroup::from_read(Arc::clone(file), path, addr)
            .into_pyobject(py)?
            .into_any()
            .unbind()),
        Some(Kind::Datatype) => Err(PyTypeError::new_err(format!(
            "{}: committed (named) datatypes are not supported by clawhdf5",
            name(&path)
        ))),
        None => Err(PyValueError::new_err(format!(
            "{}: not a dataset, group or datatype",
            name(&path)
        ))),
    }
}

/// The dataspace message of an object header.
pub(crate) fn dataspace(file: &clawhdf5_rs::File, hdr: &ObjectHeader) -> PyResult<Dataspace> {
    crate::no_panic(|| {
        let sb = file.superblock();
        let msg = hdr
            .messages
            .iter()
            .find(|m| m.msg_type == MessageType::Dataspace)
            .ok_or_else(|| PyValueError::new_err("object has no dataspace message"))?;
        let data = clawhdf5_format::shared_message::message_data(
            file.as_bytes(),
            msg,
            sb.offset_size,
            sb.length_size,
        )
        .map_err(|e| PyValueError::new_err(e.to_string()))?;
        Dataspace::parse(&data, sb.length_size).map_err(|e| PyValueError::new_err(e.to_string()))
    })
}

/// The chunk shape of a chunked dataset (one entry per dataset dimension),
/// or `None` for other layouts or a layout message that does not parse.
pub(crate) fn chunk_shape(
    file: &clawhdf5_rs::File,
    hdr: &ObjectHeader,
    rank: usize,
) -> Option<Vec<u64>> {
    let sb = file.superblock();
    let msg = hdr
        .messages
        .iter()
        .find(|m| m.msg_type == MessageType::DataLayout)?;
    match clawhdf5_format::data_layout::DataLayout::parse(&msg.data, sb.offset_size, sb.length_size)
        .ok()?
    {
        clawhdf5_format::data_layout::DataLayout::Chunked {
            chunk_dimensions, ..
        } if chunk_dimensions.len() >= rank => Some(
            chunk_dimensions[..rank]
                .iter()
                .map(|&d| u64::from(d))
                .collect(),
        ),
        _ => None,
    }
}

pub(crate) fn is_null(space: &Dataspace) -> bool {
    space.space_type == DataspaceType::Null
}

/// The attributes of the object at `addr` (whose path is `path`), sorted by
/// name (h5py's order). Attributes whose messages cannot be parsed are left
/// out, as the facade's `attrs()` does.
pub(crate) fn attributes(
    file: &clawhdf5_rs::File,
    addr: u64,
    path: &str,
) -> PyResult<Vec<AttributeMessage>> {
    let hdr = header_at(file, addr, path)?;
    crate::no_panic(|| {
        let sb = file.superblock();
        let (mut attrs, _errors) = clawhdf5_format::attribute::extract_attributes_tolerant(
            file.as_bytes(),
            &hdr,
            sb.offset_size,
            sb.length_size,
        )
        .map_err(|e| PyValueError::new_err(format!("{}: {e}", name(path))))?;
        attrs.sort_by(|a, b| a.name.as_bytes().cmp(b.name.as_bytes()));
        Ok(attrs)
    })
}

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

    #[test]
    fn join_paths() {
        assert_eq!(join("", "a"), "a");
        assert_eq!(join("a", "b/c"), "a/b/c");
        assert_eq!(join("a/b", "/x"), "x");
        assert_eq!(join("a", "/"), "");
        assert_eq!(join("", "/a//b/"), "a/b");
        assert_eq!(join("a", "./b"), "a/b");
    }
}