1#[cfg(test)]
2mod test;
3
4use crate::io::Write;
5use std::{
6 fs::File,
7 io::{self, BufReader, BufWriter, Error, ErrorKind, Read, Result},
8 mem::size_of,
9 path::Path,
10 str::from_utf8,
11};
12
13pub trait NpyType: Copy {
14 const DESCR: &'static str;
15 const SIZE: usize;
16 fn write_le(self, buffer: &mut Vec<u8>);
17 fn read_le(bytes: &[u8]) -> Self;
18 #[cfg(target_endian = "little")]
19 fn read_from<R: Read>(file: &mut R, count: usize) -> Result<Vec<Self>> {
20 let mut data = Vec::<Self>::with_capacity(count);
21 unsafe {
22 let bytes =
23 std::slice::from_raw_parts_mut(data.as_mut_ptr() as *mut u8, count * Self::SIZE);
24 file.read_exact(bytes)
25 .map_err(|_| invalid("truncated .npy data".into()))?;
26 data.set_len(count);
27 }
28 Ok(data)
29 }
30 #[cfg(target_endian = "big")]
31 fn read_from<R: Read>(file: &mut R, count: usize) -> Result<Vec<Self>> {
32 let mut bytes = vec![0u8; count * Self::SIZE];
33 file.read_exact(&mut bytes)
34 .map_err(|_| invalid("truncated .npy data".into()))?;
35 Ok(Self::read_le_all(&bytes))
36 }
37 fn read_le_all(bytes: &[u8]) -> Vec<Self> {
38 bytes.chunks_exact(Self::SIZE).map(Self::read_le).collect()
39 }
40 #[cfg(target_endian = "little")]
41 fn write_le_all<W: io::Write>(data: &[Self], file: &mut W) -> Result<()> {
42 let bytes = unsafe {
43 std::slice::from_raw_parts(data.as_ptr() as *const u8, std::mem::size_of_val(data))
44 };
45 file.write_all(bytes)
46 }
47 #[cfg(target_endian = "big")]
48 fn write_le_all<W: io::Write>(data: &[Self], file: &mut W) -> Result<()> {
49 file.write_all(&Self::write_le_bytes(data))
50 }
51 fn write_le_bytes(data: &[Self]) -> Vec<u8> {
52 let mut buffer = Vec::with_capacity(std::mem::size_of_val(data));
53 for &value in data {
54 value.write_le(&mut buffer);
55 }
56 buffer
57 }
58}
59
60macro_rules! npy_type {
61 ($type:ty, $descr:literal) => {
62 impl NpyType for $type {
63 const DESCR: &'static str = $descr;
64 const SIZE: usize = size_of::<$type>();
65 fn write_le(self, buffer: &mut Vec<u8>) {
66 buffer.extend_from_slice(&self.to_le_bytes());
67 }
68 fn read_le(bytes: &[u8]) -> Self {
69 Self::from_le_bytes(bytes.try_into().unwrap())
70 }
71 }
72 };
73}
74npy_type!(u8, "|u1");
75npy_type!(i8, "|i1");
76npy_type!(u16, "<u2");
77npy_type!(i16, "<i2");
78npy_type!(u32, "<u4");
79npy_type!(i32, "<i4");
80npy_type!(u64, "<u8");
81npy_type!(i64, "<i8");
82npy_type!(f32, "<f4");
83npy_type!(f64, "<f8");
84
85pub struct Npy<T> {
86 pub data: Vec<T>,
87 pub shape: Vec<usize>,
88 pub fortran_order: bool,
89}
90
91impl<T, P> Write<P> for Npy<T>
92where
93 T: NpyType,
94 P: AsRef<Path>,
95{
96 type Error = Error;
97 fn write(&self, path: P) -> Result<()> {
98 self.write_to(&mut BufWriter::new(File::create(path)?))
99 }
100}
101
102impl<T: NpyType> Npy<T> {
103 pub fn write_to<W: io::Write>(&self, file: &mut W) -> Result<()> {
104 write_npy(&self.data, &self.shape, self.fortran_order, file)
105 }
106}
107
108fn write_npy<T: NpyType, W: io::Write>(
109 data: &[T],
110 shape: &[usize],
111 fortran_order: bool,
112 file: &mut W,
113) -> Result<()> {
114 let order = if fortran_order { "True" } else { "False" };
115 let dims: String = shape.iter().map(|d| format!("{d}, ")).collect();
116 let mut header = format!(
117 "{{'descr': '{}', 'fortran_order': {order}, 'shape': ({dims}), }}",
118 T::DESCR
119 );
120 let pad = (64 - (10 + header.len() + 1) % 64) % 64;
121 header.push_str(&" ".repeat(pad));
122 header.push('\n');
123 file.write_all(b"\x93NUMPY")?;
124 file.write_all(&[1, 0])?;
125 file.write_all(&(header.len() as u16).to_le_bytes())?;
126 file.write_all(header.as_bytes())?;
127 T::write_le_all(data, file)?;
128 file.flush()
129}
130
131impl<T: NpyType> Npy<T> {
132 pub fn read<P: AsRef<Path>>(path: P) -> Result<Self> {
133 let mut file = BufReader::new(File::open(path)?);
134 Self::read_from(&mut file)
135 }
136 pub fn read_from<R: Read>(file: &mut R) -> Result<Self> {
137 let mut prefix = [0u8; 10];
138 file.read_exact(&mut prefix)
139 .map_err(|_| invalid("not a .npy file".into()))?;
140 if &prefix[..6] != b"\x93NUMPY" {
141 return Err(invalid("not a .npy file".into()));
142 }
143 let header_length = match prefix[6] {
144 1 => u16::from_le_bytes([prefix[8], prefix[9]]) as usize,
145 2 => {
146 let mut rest = [0u8; 2];
147 file.read_exact(&mut rest)?;
148 u32::from_le_bytes([prefix[8], prefix[9], rest[0], rest[1]]) as usize
149 }
150 other => return Err(invalid(format!("unsupported .npy version {other}"))),
151 };
152 let mut header_bytes = vec![0u8; header_length];
153 file.read_exact(&mut header_bytes)?;
154 let header =
155 from_utf8(&header_bytes).map_err(|_| invalid("non-UTF-8 .npy header".into()))?;
156 let descr = quoted(header, "'descr':").ok_or_else(|| invalid("no descr".into()))?;
157 if descr.starts_with('>') || descr.get(1..) != T::DESCR.get(1..) {
158 return Err(invalid(format!(
159 "dtype {descr} does not match {}",
160 T::DESCR
161 )));
162 }
163 let fortran_order = header.contains("'fortran_order': True");
164 let shape = shape(header)?;
165 let count: usize = shape.iter().product();
166 let data = T::read_from(file, count)?;
167 Ok(Npy {
168 data,
169 shape,
170 fortran_order,
171 })
172 }
173}
174
175fn quoted<'a>(header: &'a str, key: &str) -> Option<&'a str> {
176 let at = header.find(key)? + key.len();
177 let open = header[at..].find('\'')? + at + 1;
178 let close = header[open..].find('\'')? + open;
179 Some(&header[open..close])
180}
181
182fn shape(header: &str) -> Result<Vec<usize>> {
183 let at = header
184 .find("'shape':")
185 .ok_or_else(|| invalid("no shape".into()))?;
186 let open = header[at..]
187 .find('(')
188 .ok_or_else(|| invalid("malformed shape".into()))?
189 + at
190 + 1;
191 let close = header[open..]
192 .find(')')
193 .ok_or_else(|| invalid("malformed shape".into()))?
194 + open;
195 header[open..close]
196 .split(',')
197 .map(str::trim)
198 .filter(|dim| !dim.is_empty())
199 .map(|dim| {
200 dim.parse()
201 .map_err(|_| invalid(format!("bad shape entry {dim}")))
202 })
203 .collect()
204}
205
206fn invalid(message: String) -> Error {
207 Error::new(ErrorKind::InvalidData, message)
208}