clawhdf5/crates/clawhdf5/tests/userblock_interop.rs

//! Files that start with a user block (`h5py.File(..., userblock_size=N)`,
//! `h5jam`): the superblock sits at 512, 1024, ... and every address in the
//! file is relative to it. Each reader (buffered, mmap, `MmapFile`,
//! `LazyFile`) must apply that base, and read the same values h5py does.
//!
//! h5py writes the files; skipped when python3 with h5py is unavailable,
//! unless `CLAWHDF5_REQUIRE_INTEROP=1`.

use std::collections::HashMap;
use std::path::Path;
use std::process::Command;

use clawhdf5::{AttrValue, File, LazyFile, MmapFile};

fn python() -> String {
    std::env::var("CLAWHDF5_PYTHON").unwrap_or_else(|_| "python3".to_string())
}

fn interop_required() -> bool {
    std::env::var("CLAWHDF5_REQUIRE_INTEROP").is_ok_and(|v| v == "1")
}

fn python_available() -> bool {
    Command::new(python())
        .args(["-c", "import h5py, numpy"])
        .output()
        .map(|o| o.status.success())
        .unwrap_or(false)
}

macro_rules! skip_if_no_python {
    () => {
        if !python_available() {
            assert!(
                !interop_required(),
                "CLAWHDF5_REQUIRE_INTEROP=1 but python3 with h5py is not available"
            );
            eprintln!("SKIP: python3 with h5py not available");
            return;
        }
    };
}

/// Run `script` and return its stdout as `key -> values` (one
/// `key v1 v2 ...` line per key).
fn run_python(script: &str) -> HashMap<String, Vec<String>> {
    let output = Command::new(python())
        .args(["-c", script])
        .output()
        .expect("failed to run python");
    assert!(
        output.status.success(),
        "python failed:\n{}",
        String::from_utf8_lossy(&output.stderr)
    );
    String::from_utf8_lossy(&output.stdout)
        .lines()
        .filter_map(|line| {
            let mut words = line.split_whitespace().map(str::to_string);
            Some((words.next()?, words.collect()))
        })
        .collect()
}

fn parse<T: std::str::FromStr>(values: &[String]) -> Vec<T>
where
    T::Err: std::fmt::Debug,
{
    values.iter().map(|v| v.parse().unwrap()).collect()
}

/// Write a file with a user block of `userblock` bytes holding contiguous,
/// chunked (deflate), compact and committed-type datasets, nested groups,
/// and attributes (compact and, under `latest`, dense). Prints what h5py
/// reads back.
fn write_file(path: &Path, userblock: u32, libver: &str) -> HashMap<String, Vec<String>> {
    let script = format!(
        r#"
import h5py, numpy as np
path = "{path}"
with h5py.File(path, "w", userblock_size={userblock}, libver={libver}) as f:
    f.attrs["title"] = "user block"
    f.attrs["answer"] = np.int64(42)
    f.create_dataset("contig", data=np.arange(12, dtype="<f8") * 0.5)
    f.create_dataset("chunked", data=np.arange(1000, dtype="<i4") * 3 - 7,
                     chunks=(128,), compression="gzip")
    dcpl = h5py.h5p.create(h5py.h5p.DATASET_CREATE)
    dcpl.set_layout(h5py.h5d.COMPACT)
    space = h5py.h5s.create_simple((5,))
    dsid = h5py.h5d.create(f.id, b"compact", h5py.h5t.STD_I64LE, space, dcpl=dcpl)
    dsid.write(h5py.h5s.ALL, h5py.h5s.ALL, np.array([5, -4, 3, -2, 1], "<i8"))
    f["named_type"] = np.dtype("<f4")
    f.create_dataset("committed", data=np.array([1.25, -2.5], "<f4"),
                     dtype=f["named_type"])
    g = f.create_group("a/b")
    g.create_dataset("deep", data=np.array([7, 8, 9], "<i8"))
    g.attrs["scale"] = 2.5
    d = f["contig"]
    d.attrs["units"] = "m"
    # Enough attributes that `latest` stores them densely (fractal heap).
    for i in range(12):
        f["a"].attrs["k%02d" % i] = np.int64(i * i)
    # And enough links for a dense (fractal-heap) group under `latest`.
    many = f.create_group("many")
    for i in range(20):
        many.create_dataset("d%02d" % i, data=np.array([i], "<i4"))
with h5py.File(path, "r") as f:
    print("userblock", f.userblock_size)
    print("contig", *f["contig"][()])
    print("chunked", *f["chunked"][()])
    print("compact", *f["compact"][()])
    print("committed", *f["committed"][()])
    print("deep", *f["a/b/deep"][()])
    print("many", *[int(f["many/d%02d" % i][0]) for i in range(20)])
    print("k", *[int(f["a"].attrs["k%02d" % i]) for i in range(12)])
"#,
        path = path.display(),
        libver = if libver == "default" {
            "None".to_string()
        } else {
            format!("{libver:?}")
        },
    );
    run_python(&script)
}

/// Attribute value rendered for comparison (`AttrValue` has no `PartialEq`).
fn attr(map: &HashMap<String, AttrValue>, key: &str) -> String {
    match map.get(key) {
        Some(AttrValue::I64(v)) => format!("i64 {v}"),
        Some(AttrValue::F64(v)) => format!("f64 {v}"),
        Some(AttrValue::String(v)) => format!("str {v}"),
        other => format!("{other:?}"),
    }
}

fn i64s(v: &[String]) -> Vec<i64> {
    parse(v)
}

/// Everything read through the `File` API must match h5py.
fn check_file(file: &File, expected: &HashMap<String, Vec<String>>, label: &str) {
    let ub: u64 = expected["userblock"][0].parse().unwrap();
    assert_eq!(file.user_block_size(), ub, "{label}: user block size");
    assert_eq!(
        file.dataset("contig").unwrap().read_f64().unwrap(),
        parse::<f64>(&expected["contig"]),
        "{label}: contiguous"
    );
    assert_eq!(
        file.dataset("chunked")
            .unwrap()
            .read_i32()
            .unwrap()
            .iter()
            .map(|&v| v as i64)
            .collect::<Vec<_>>(),
        i64s(&expected["chunked"]),
        "{label}: chunked"
    );
    assert_eq!(
        file.dataset("compact").unwrap().read_i64().unwrap(),
        i64s(&expected["compact"]),
        "{label}: compact"
    );
    assert_eq!(
        file.dataset("committed").unwrap().read_f32().unwrap(),
        parse::<f32>(&expected["committed"]),
        "{label}: committed datatype"
    );
    assert_eq!(
        file.dataset("a/b/deep").unwrap().read_i64().unwrap(),
        i64s(&expected["deep"]),
        "{label}: nested group"
    );
    let many: Vec<i64> = (0..20)
        .map(|i| {
            file.dataset(&format!("many/d{i:02}"))
                .unwrap()
                .read_i32()
                .unwrap()[0] as i64
        })
        .collect();
    assert_eq!(many, i64s(&expected["many"]), "{label}: many links");

    let root = file.root().attrs().unwrap();
    assert_eq!(attr(&root, "title"), "str user block", "{label}");
    assert_eq!(attr(&root, "answer"), "i64 42", "{label}");
    let a = file.group("a").unwrap().attrs().unwrap();
    let k: Vec<i64> = (0..12)
        .map(|i| match &a[&format!("k{i:02}")] {
            AttrValue::I64(v) => *v,
            _ => panic!("{label}: k{i:02} is not an i64"),
        })
        .collect();
    assert_eq!(k, i64s(&expected["k"]), "{label}: attributes");
    assert_eq!(
        attr(&file.group("a/b").unwrap().attrs().unwrap(), "scale"),
        "f64 2.5",
        "{label}"
    );
    assert_eq!(
        attr(&file.dataset("contig").unwrap().attrs().unwrap(), "units"),
        "str m",
        "{label}"
    );
}

fn check_all_readers(path: &Path, expected: &HashMap<String, Vec<String>>, label: &str) {
    check_file(
        &File::open(path).unwrap(),
        expected,
        &format!("{label} File::open"),
    );
    check_file(
        &File::open_buffered(path).unwrap(),
        expected,
        &format!("{label} File::open_buffered"),
    );
    check_file(
        &File::from_bytes(std::fs::read(path).unwrap()).unwrap(),
        expected,
        &format!("{label} File::from_bytes"),
    );

    let ub: u64 = expected["userblock"][0].parse().unwrap();

    let mm = MmapFile::open(path).unwrap();
    assert_eq!(mm.user_block_size(), ub, "{label} MmapFile");
    assert_eq!(
        mm.dataset("contig").unwrap().read_f64().unwrap(),
        parse::<f64>(&expected["contig"]),
        "{label} MmapFile contiguous"
    );
    assert_eq!(
        mm.dataset("compact").unwrap().read_i64().unwrap(),
        i64s(&expected["compact"]),
        "{label} MmapFile compact"
    );
    assert_eq!(
        mm.dataset("committed").unwrap().read_f32().unwrap(),
        parse::<f32>(&expected["committed"]),
        "{label} MmapFile committed"
    );
    assert_eq!(
        mm.dataset("a/b/deep").unwrap().read_i64().unwrap(),
        i64s(&expected["deep"]),
        "{label} MmapFile nested"
    );
    assert_eq!(
        attr(&mm.root().attrs().unwrap(), "answer"),
        "i64 42",
        "{label} MmapFile attrs"
    );

    let lazy = LazyFile::open_mmap(path).unwrap();
    assert_eq!(lazy.user_block_size(), ub, "{label} LazyFile");
    assert_eq!(
        lazy.dataset("contig").unwrap().read_f64().unwrap(),
        parse::<f64>(&expected["contig"]),
        "{label} LazyFile contiguous"
    );
    assert_eq!(
        lazy.dataset("chunked")
            .unwrap()
            .read_i32()
            .unwrap()
            .iter()
            .map(|&v| v as i64)
            .collect::<Vec<_>>(),
        i64s(&expected["chunked"]),
        "{label} LazyFile chunked"
    );
    assert_eq!(
        lazy.dataset("committed").unwrap().read_f32().unwrap(),
        parse::<f32>(&expected["committed"]),
        "{label} LazyFile committed"
    );
    assert_eq!(
        lazy.dataset("a/b/deep").unwrap().read_i64().unwrap(),
        i64s(&expected["deep"]),
        "{label} LazyFile nested"
    );
    assert_eq!(
        attr(&lazy.root().attrs().unwrap(), "answer"),
        "i64 42",
        "{label} LazyFile attrs"
    );
}

#[test]
fn user_block_files_read_like_h5py() {
    skip_if_no_python!();
    let dir = tempfile::tempdir().unwrap();
    for userblock in [512u32, 4096] {
        for libver in ["default", "latest"] {
            let label = format!("userblock={userblock} libver={libver}");
            let path = dir.path().join(format!("ub_{userblock}_{libver}.h5"));
            let expected = write_file(&path, userblock, libver);
            assert_eq!(expected["userblock"], [userblock.to_string()], "{label}");
            check_all_readers(&path, &expected, &label);
        }
    }
}

#[test]
fn file_without_user_block_reports_zero() {
    skip_if_no_python!();
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("no_ub.h5");
    let expected = write_file(&path, 0, "default");
    assert_eq!(expected["userblock"], ["0"]);
    check_all_readers(&path, &expected, "userblock=0");
}