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

//! Resolving paths to objects in a file opened for reading.
//!
//! Everything here parses through `File::storage()` (the `clawhdf5_format`
//! `*_in` functions), never `File::as_bytes()`, so it works the same on a
//! memory-mapped local file and on a remote one; and it runs inside
//! `Handle::with`, without the GIL.

use std::sync::Arc;

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

use crate::dataset::{DatasetMeta, PyDataset};
use crate::group::PyGroup;
use crate::handle::Handle;

/// 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}")
}

/// A format error met at `path`: a failed read of the storage (a network
/// error on a remote file) is an `OSError`, anything else `other(message)`.
pub(crate) fn format_err(path: &str, e: FormatError, other: fn(String) -> PyErr) -> PyErr {
    let msg = format!("{}: {e}", name(path));
    match e {
        FormatError::Storage(_) => PyOSError::new_err(msg),
        _ => other(msg),
    }
}

fn value_err(msg: String) -> PyErr {
    PyValueError::new_err(msg)
}

/// The address of the object at `path`, resolved from the root group.
pub(crate) fn address(file: &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: &File, group: u64, rel: &str, full: &str) -> PyResult<u64> {
    if rel.is_empty() {
        return Ok(group);
    }
    clawhdf5_format::group_v2::resolve_path_from_in(file.storage(), file.superblock(), group, rel)
        .map_err(|e| match e {
            FormatError::Storage(_) => format_err(full, e, value_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: &File, addr: u64, path: &str) -> PyResult<ObjectHeader> {
    let sb = file.superblock();
    ObjectHeader::parse_in(file.storage(), addr, sb.offset_size, sb.length_size)
        .map_err(|e| format_err(path, e, value_err))
}

/// 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
    }
}

/// The kind of the object at `addr`, from its header.
pub(crate) fn kind_at(file: &File, addr: u64, path: &str) -> PyResult<Option<Kind>> {
    Ok(kind(&header_at(file, addr, path)?))
}

/// What opening an object found, read without the GIL.
enum Found {
    Dataset(DatasetMeta),
    Group,
    Datatype,
    Other,
}

/// 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<'_>,
    handle: &Arc<Handle>,
    path: String,
    addr: u64,
) -> PyResult<Py<PyAny>> {
    let found = handle.with(py, |f| {
        let hdr = header_at(f, addr, &path)?;
        Ok(match kind(&hdr) {
            Some(Kind::Dataset) => Found::Dataset(DatasetMeta::load(f, addr, &hdr, &path)?),
            Some(Kind::Group) => Found::Group,
            Some(Kind::Datatype) => Found::Datatype,
            None => Found::Other,
        })
    })?;
    match found {
        Found::Dataset(meta) => Ok(PyDataset::new(py, Arc::clone(handle), path, addr, meta)
            .into_pyobject(py)?
            .into_any()
            .unbind()),
        Found::Group => Ok(PyGroup::from_read(Arc::clone(handle), path, addr)
            .into_pyobject(py)?
            .into_any()
            .unbind()),
        Found::Datatype => Err(PyTypeError::new_err(format!(
            "{}: committed (named) datatypes are not supported by clawhdf5",
            name(&path)
        ))),
        Found::Other => 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: &File, hdr: &ObjectHeader, path: &str) -> PyResult<Dataspace> {
    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_in(
        file.storage(),
        msg,
        sb.offset_size,
        sb.length_size,
    )
    .map_err(|e| format_err(path, e, value_err))?;
    Dataspace::parse(&data, sb.length_size).map_err(|e| format_err(path, e, value_err))
}

/// 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: &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: &File, addr: u64, path: &str) -> PyResult<Vec<AttributeMessage>> {
    let hdr = header_at(file, addr, path)?;
    let sb = file.superblock();
    let (mut attrs, _errors) = clawhdf5_format::attribute::extract_attributes_tolerant_in(
        file.storage(),
        &hdr,
        sb.offset_size,
        sb.length_size,
    )
    .map_err(|e| format_err(path, e, value_err))?;
    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");
    }
}