nuvec_io/
npy.rs

1//! NumPy's `.npy` format: a magic string, a version, a Python-literal header (`descr`,
2//! `fortran_order`, `shape`), then the elements.
3//!
4//! Loading maps the file without a copy (see [`MappedFile`]); Fortran-order files load as
5//! transposed views.
6
7use std::path::Path;
8
9use nuvec_core::{Array, DType, Error, MappedFile, Result};
10
11use crate::safetensors::array_bytes;
12
13fn invalid(msg: impl Into<String>) -> Error {
14    Error::InvalidArgument(format!("npy: {}", msg.into()))
15}
16
17/// NumPy's type string (little-endian) for a dtype; bfloat16 has none.
18fn descr_of(dtype: DType) -> Option<&'static str> {
19    Some(match dtype {
20        DType::Bool => "|b1",
21        DType::U8 => "|u1",
22        DType::U16 => "<u2",
23        DType::U32 => "<u4",
24        DType::U64 => "<u8",
25        DType::I8 => "|i1",
26        DType::I16 => "<i2",
27        DType::I32 => "<i4",
28        DType::I64 => "<i8",
29        DType::F16 => "<f2",
30        DType::F32 => "<f4",
31        DType::F64 => "<f8",
32        DType::Complex64 => "<c8",
33        _ => return None,
34    })
35}
36
37fn dtype_of(descr: &str) -> Option<DType> {
38    // Single-byte types carry no byte order; `=` is native (little-endian on Apple Silicon).
39    let canonical = match descr.as_bytes() {
40        [b'=' | b'<' | b'|', rest @ ..] => format!("<{}", std::str::from_utf8(rest).ok()?),
41        _ => return None,
42    };
43    [
44        DType::Bool,
45        DType::U8,
46        DType::U16,
47        DType::U32,
48        DType::U64,
49        DType::I8,
50        DType::I16,
51        DType::I32,
52        DType::I64,
53        DType::F16,
54        DType::F32,
55        DType::F64,
56        DType::Complex64,
57    ]
58    .into_iter()
59    .find(|&d| descr_of(d).is_some_and(|s| s[1..] == canonical[1..]))
60}
61
62/// The text after `'key':` in the header dict.
63fn field<'a>(header: &'a str, key: &str) -> Result<&'a str> {
64    let pat = format!("'{key}':");
65    let at = header
66        .find(&pat)
67        .ok_or_else(|| invalid(format!("header has no `{key}`")))?;
68    Ok(header[at + pat.len()..].trim_start())
69}
70
71/// Loads a `.npy` file as a view of the memory-mapped file.
72pub fn load_npy(path: impl AsRef<Path>) -> Result<Array> {
73    let file = MappedFile::open(path)?;
74    let bytes = file.bytes();
75    if bytes.get(..6) != Some(b"\x93NUMPY") || bytes.len() < 10 {
76        return Err(invalid("not an npy file"));
77    }
78    let (len, start) = match bytes[6] {
79        1 => (usize::from(u16::from_le_bytes([bytes[8], bytes[9]])), 10),
80        2 | 3 => (
81            bytes
82                .get(8..12)
83                .map(|b| u32::from_le_bytes(b.try_into().expect("4 bytes")) as usize)
84                .ok_or_else(|| invalid("truncated header"))?,
85            12,
86        ),
87        v => return Err(invalid(format!("unsupported version {v}"))),
88    };
89    let header = bytes
90        .get(start..start + len)
91        .and_then(|h| std::str::from_utf8(h).ok())
92        .ok_or_else(|| invalid("truncated or non-text header"))?;
93    let descr = field(header, "descr")?;
94    let descr = descr
95        .strip_prefix('\'')
96        .and_then(|d| d.split('\'').next())
97        .ok_or_else(|| invalid("bad descr"))?;
98    let dtype = dtype_of(descr).ok_or_else(|| invalid(format!("unsupported type {descr}")))?;
99    let fortran = field(header, "fortran_order")?.starts_with("True");
100    // `(d0, d1, ...)`; sliced by pattern, never by byte index (the header may be any UTF-8).
101    let (dims, _) = field(header, "shape")?
102        .strip_prefix('(')
103        .and_then(|s| s.split_once(')'))
104        .ok_or_else(|| invalid("bad shape"))?;
105    let shape = dims
106        .split(',')
107        .map(str::trim)
108        .filter(|s| !s.is_empty())
109        .map(|s| s.parse::<usize>().map_err(|_| invalid("bad shape")))
110        .collect::<Result<Vec<_>>>()?;
111    let data = start + len;
112    if fortran && shape.len() > 1 {
113        let reversed: Vec<usize> = shape.iter().rev().copied().collect();
114        Ok(file.array(data, dtype, &reversed)?.t())
115    } else {
116        file.array(data, dtype, &shape)
117    }
118}
119
120/// Saves `a` (evaluating it) as a C-order `.npy` file (version 1). bfloat16 has no NumPy
121/// type and is rejected.
122pub fn save_npy(path: impl AsRef<Path>, a: &Array) -> Result<()> {
123    let descr =
124        descr_of(a.dtype()).ok_or_else(|| invalid(format!("{} has no NumPy type", a.dtype())))?;
125    let shape = match a.shape() {
126        [n] => format!("({n},)"),
127        s => format!(
128            "({})",
129            s.iter().map(i32::to_string).collect::<Vec<_>>().join(", ")
130        ),
131    };
132    let mut header = format!("{{'descr': '{descr}', 'fortran_order': False, 'shape': {shape}, }}");
133    // Pad so the data starts 64-byte aligned; the header ends with a newline.
134    while (10 + header.len() + 1) % 64 != 0 {
135        header.push(' ');
136    }
137    header.push('\n');
138    let mut out = b"\x93NUMPY\x01\x00".to_vec();
139    out.extend((header.len() as u16).to_le_bytes());
140    out.extend(header.into_bytes());
141    out.extend(array_bytes(a)?);
142    std::fs::write(path.as_ref(), out)
143        .map_err(|e| Error::Io(format!("{}: {e}", path.as_ref().display())))
144}
nuvec_io/npy.rs

nuvec_io/
npy.rs
