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, transpose};
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    let shape_text = field(header, "shape")?;
101    let close = shape_text.find(')').ok_or_else(|| invalid("bad shape"))?;
102    let shape = shape_text[1..close]
103        .split(',')
104        .map(str::trim)
105        .filter(|s| !s.is_empty())
106        .map(|s| s.parse::<usize>().map_err(|_| invalid("bad shape")))
107        .collect::<Result<Vec<_>>>()?;
108    let data = start + len;
109    if fortran && shape.len() > 1 {
110        let reversed: Vec<usize> = shape.iter().rev().copied().collect();
111        let a = file.array(data, dtype, &reversed)?;
112        let axes: Vec<usize> = (0..shape.len()).rev().collect();
113        transpose(&a, &axes)
114    } else {
115        file.array(data, dtype, &shape)
116    }
117}
118
119/// Saves `a` (evaluating it) as a C-order `.npy` file (version 1). bfloat16 has no NumPy
120/// type and is rejected.
121pub fn save_npy(path: impl AsRef<Path>, a: &Array) -> Result<()> {
122    let descr =
123        descr_of(a.dtype()).ok_or_else(|| invalid(format!("{} has no NumPy type", a.dtype())))?;
124    let shape = match a.shape() {
125        [n] => format!("({n},)"),
126        s => format!(
127            "({})",
128            s.iter()
129                .map(usize::to_string)
130                .collect::<Vec<_>>()
131                .join(", ")
132        ),
133    };
134    let mut header = format!("{{'descr': '{descr}', 'fortran_order': False, 'shape': {shape}, }}");
135    // Pad so the data starts 64-byte aligned; the header ends with a newline.
136    while (10 + header.len() + 1) % 64 != 0 {
137        header.push(' ');
138    }
139    header.push('\n');
140    let mut out = b"\x93NUMPY\x01\x00".to_vec();
141    out.extend((header.len() as u16).to_le_bytes());
142    out.extend(header.into_bytes());
143    out.extend(array_bytes(a)?);
144    std::fs::write(path.as_ref(), out)
145        .map_err(|e| Error::Io(format!("{}: {e}", path.as_ref().display())))
146}
nuvec_io/npy.rs

nuvec_io/
npy.rs
