sparse_npz 0.1.0

//! Minimal reader and writer for NumPy `.npy` arrays, covering the pieces of a
//! SciPy sparse `.npz` bundle: little-endian 1-D integer/float arrays and the
//! zero-dimensional `|S3` "format" string. Only C-order, little-endian data is
//! handled, which is what NumPy and SciPy emit on the platforms we target.

const MAGIC: &[u8] = b"\x93NUMPY";

/// A parsed `.npy` array: its NumPy dtype string and the raw little-endian
/// element bytes. The declared shape is validated during parsing but not
/// retained, since callers here read flat 1-D arrays and the scalar format tag.
pub struct NpyArray {
    pub descr: String,
    pub data: Vec<u8>,
}

impl NpyArray {
    /// Parses an in-memory `.npy` file.
    pub fn parse(bytes: &[u8]) -> Result<NpyArray, String> {
        if bytes.len() < 10 || &bytes[..6] != MAGIC {
            return Err("not a .npy array".into());
        }
        let major = bytes[6];
        let (header, data_start) = if major == 1 {
            let len = u16::from_le_bytes([bytes[8], bytes[9]]) as usize;
            (&bytes[10..10 + len], 10 + len)
        } else {
            let len = u32::from_le_bytes([bytes[8], bytes[9], bytes[10], bytes[11]]) as usize;
            (&bytes[12..12 + len], 12 + len)
        };
        let header = std::str::from_utf8(header).map_err(|_| "non-UTF8 npy header")?;

        let descr = extract_quoted(header, "'descr'").ok_or("npy header lacks descr")?;
        if descr.starts_with('>') {
            return Err("big-endian .npy arrays are not supported".into());
        }
        extract_shape(header).ok_or("npy header lacks shape")?;
        Ok(NpyArray {
            descr,
            data: bytes[data_start..].to_vec(),
        })
    }

    /// Interprets the array's elements as `i64`, widening any integer dtype.
    pub fn as_i64(&self) -> Result<Vec<i64>, String> {
        Ok(self.as_f64()?.into_iter().map(|v| v as i64).collect())
    }

    /// Interprets the array's elements as `f64`, widening any integer, boolean,
    /// or float dtype.
    pub fn as_f64(&self) -> Result<Vec<f64>, String> {
        let d = &self.data;
        Ok(match self.descr.as_str() {
            "|b1" | "<b1" => d.iter().map(|&b| if b != 0 { 1.0 } else { 0.0 }).collect(),
            "<i1" | "|i1" => d.iter().map(|&b| b as i8 as f64).collect(),
            "<u1" | "|u1" => d.iter().map(|&b| b as f64).collect(),
            "<i2" => map_chunks::<2>(d, |a| i16::from_le_bytes(a) as f64)?,
            "<u2" => map_chunks::<2>(d, |a| u16::from_le_bytes(a) as f64)?,
            "<i4" => map_chunks::<4>(d, |a| i32::from_le_bytes(a) as f64)?,
            "<u4" => map_chunks::<4>(d, |a| u32::from_le_bytes(a) as f64)?,
            "<i8" => map_chunks::<8>(d, |a| i64::from_le_bytes(a) as f64)?,
            "<u8" => map_chunks::<8>(d, |a| u64::from_le_bytes(a) as f64)?,
            "<f4" => map_chunks::<4>(d, |a| f32::from_le_bytes(a) as f64)?,
            "<f8" => map_chunks::<8>(d, f64::from_le_bytes)?,
            other => return Err(format!("unsupported numeric npy dtype {other}")),
        })
    }

    /// Returns the raw element bytes with trailing NULs trimmed, for reading the
    /// `|S`-typed `format` string.
    pub fn as_ascii(&self) -> String {
        let end = self.data.iter().rposition(|&b| b != 0).map_or(0, |p| p + 1);
        String::from_utf8_lossy(&self.data[..end]).into_owned()
    }
}

fn map_chunks<const N: usize>(data: &[u8], f: impl Fn([u8; N]) -> f64) -> Result<Vec<f64>, String> {
    if data.len() % N != 0 {
        return Err("npy data length not a multiple of element size".into());
    }
    Ok(data
        .chunks_exact(N)
        .map(|c| {
            let mut a = [0u8; N];
            a.copy_from_slice(c);
            f(a)
        })
        .collect())
}

fn extract_quoted(header: &str, key: &str) -> Option<String> {
    let after = &header[header.find(key)? + key.len()..];
    let after = &after[after.find(':')? + 1..];
    let start = after.find('\'')? + 1;
    let rest = &after[start..];
    let end = rest.find('\'')?;
    Some(rest[..end].to_string())
}

fn extract_shape(header: &str) -> Option<Vec<usize>> {
    let after = &header[header.find("'shape'")? + 7..];
    let open = after.find('(')?;
    let close = after[open..].find(')')? + open;
    let inside = &after[open + 1..close];
    Some(
        inside
            .split(',')
            .filter_map(|s| s.trim().parse::<usize>().ok())
            .collect(),
    )
}

/// Serializes a `.npy` array from its dtype, shape, and raw little-endian bytes.
pub fn write(descr: &str, shape: &[usize], data: &[u8]) -> Vec<u8> {
    let shape_field = match shape {
        [] => "()".to_string(),
        [n] => format!("({n},)"),
        dims => format!("({})", dims.iter().map(|d| d.to_string()).collect::<Vec<_>>().join(", ")),
    };
    let mut header = format!("{{'descr': '{descr}', 'fortran_order': False, 'shape': {shape_field}, }}");

    let unpadded = 10 + header.len() + 1;
    let padding = (64 - unpadded % 64) % 64;
    header.push_str(&" ".repeat(padding));
    header.push('\n');

    let mut out = Vec::with_capacity(10 + header.len() + data.len());
    out.extend_from_slice(MAGIC);
    out.push(1);
    out.push(0);
    out.extend_from_slice(&(header.len() as u16).to_le_bytes());
    out.extend_from_slice(header.as_bytes());
    out.extend_from_slice(data);
    out
}