Skip to main content

datui_lib/formats/
numpy.rs

1//! NumPy arrays: `.npy` files, and `.npz` archives of them.
2//!
3//! An `.npy` file is a short header, a Python dict literal giving the element type
4//! (`descr`), the order (`fortran_order`) and the `shape`, and then the elements packed
5//! end to end. Every element is the same size, so the data is read where it is, through
6//! the strided decoder of the binary format specs ([`crate::formats::fixed_records`]): a map of
7//! the file, decoded only where it is shown.
8//!
9//! A structured element type gives a column per field (nested fields as `outer.inner`,
10//! a subarray field as an Array column); a 1-D array is one column, named for the
11//! array; a 2-D array is a grid with a column per index, `0` to `n-1`. More dimensions
12//! are refused with the shape. Python objects (`O`), which only unpickling could read,
13//! are refused and never unpickled.
14//!
15//! An `.npz` archive is a zip of `.npy` files. Its arrays are listed on the home screen
16//! like a directory's files; an archive of one opens it. An array stored as it is in the
17//! archive (`np.savez`) is read in place from a map of the archive; a compressed one
18//! (`np.savez_compressed`) is decompressed once to a temporary file first.
19
20use std::io::Read;
21use std::path::Path;
22use std::sync::Arc;
23
24use color_eyre::Result;
25use color_eyre::eyre::eyre;
26
27use crate::error_display::{FileError, in_file};
28use polars::prelude::*;
29
30use crate::formats::fixed_records::{Bytes, ColumnLayout, FixedRecords, Logical, Null, Physical};
31use crate::formats::model_files::MetaValue;
32use crate::formats::sqlite::Table;
33use crate::formats::text_formats::Detail;
34
35/// What datui does with a NumPy file: see [`crate::formats::readers`].
36pub(crate) const READER: crate::formats::readers::Reader = crate::formats::readers::Reader {
37    // An archive's compressed array, named by `--table` or found alone.
38    convert: Some(|input| {
39        let ([file], Some(name)) = (input.files, input.options.table.as_deref()) else {
40            return Err(eyre!("Open one array of an archive at a time."));
41        };
42        let (held, lf, opened) = convert(file, name, input.options, input.writer, input.read)?;
43        Ok((
44            crate::formats::segments::Converted {
45                lf,
46                files: vec![held],
47                notes: opened.notes,
48                other_tables: opened.other_tables,
49            },
50            opened.detail,
51        ))
52    }),
53    scan,
54    python: Some(crate::export::python_script::Python {
55        call: "pl.from_numpy",
56        eager: true,
57        glob_flag: false,
58        arguments: Some(crate::export::python_script::numpy_arguments),
59    }),
60    signatures: &[
61        crate::formats::readers::Signature {
62            says: |head, _| looks_like(head),
63            kind: crate::formats::readers::Kind::Magic,
64            trusted: crate::formats::readers::EVERYWHERE,
65        },
66        // An archive is a zip file named `.npz`: a zip file is many things.
67        crate::formats::readers::Signature {
68            says: |head, file| file.is_some_and(|file| is_archive(file, head)),
69            kind: crate::formats::readers::Kind::Magic,
70            trusted: crate::formats::readers::Trusted {
71                pipe: false,
72                open: crate::formats::readers::Unnamed::Any,
73                listing: false,
74                tables: true,
75            },
76        },
77    ],
78    tables: Some(tables),
79    table_schema: Some(schema_preview),
80    ..crate::formats::readers::BASE
81};
82
83/// The first six bytes of every `.npy` file.
84pub const MAGIC: &[u8; 6] = b"\x93NUMPY";
85
86/// Nesting of the header's literal, and of structured types inside structured types.
87const MAX_DEPTH: usize = 32;
88
89/// Columns one array may make, nested fields and repeated structs flattened.
90pub const MAX_COLUMNS: usize = 4096;
91
92/// A 2-D array of more columns than this is one Array column rather than a column per
93/// index.
94pub const MAX_GRID: usize = 1024;
95
96/// Arrays of an archive listed or read.
97const MAX_MEMBERS: usize = 100_000;
98
99/// Whether `head`, the first bytes of a file, begins a `.npy` file.
100pub fn looks_like(head: &[u8]) -> bool {
101    head.starts_with(MAGIC)
102}
103
104/// Whether a file is a NumPy archive: a zip file named `.npz`.
105pub fn is_archive(path: &Path, head: &[u8]) -> bool {
106    is_archive_name(path) && (head.starts_with(b"PK\x03\x04") || head.starts_with(b"PK\x05\x06"))
107}
108
109fn is_archive_name(path: &Path) -> bool {
110    path.extension()
111        .is_some_and(|e| e.eq_ignore_ascii_case("npz"))
112}
113
114// --- The header's Python literal -------------------------------------------------
115
116/// A Python literal, as far as a `.npy` header uses them.
117#[derive(Debug, Clone, PartialEq)]
118pub enum Literal {
119    Str(String),
120    Int(i128),
121    Bool(bool),
122    None,
123    List(Vec<Literal>),
124    Tuple(Vec<Literal>),
125    Dict(Vec<(Literal, Literal)>),
126}
127
128impl Literal {
129    fn as_str(&self) -> Option<&str> {
130        match self {
131            Self::Str(s) => Some(s),
132            _ => None,
133        }
134    }
135
136    fn items(&self) -> Option<&[Literal]> {
137        match self {
138            Self::List(items) | Self::Tuple(items) => Some(items),
139            _ => None,
140        }
141    }
142
143    fn get(&self, key: &str) -> Option<&Literal> {
144        match self {
145            Self::Dict(pairs) => pairs
146                .iter()
147                .find(|(k, _)| k.as_str() == Some(key))
148                .map(|(_, v)| v),
149            _ => None,
150        }
151    }
152}
153
154/// Parse the Python literal `text`: strings, integers, `True`, `False`, `None`, lists,
155/// tuples and dicts, nested at most `MAX_DEPTH` deep.
156pub fn parse_literal(text: &str) -> std::result::Result<Literal, String> {
157    let mut parser = Parser {
158        text: text.as_bytes(),
159        at: 0,
160    };
161    let value = parser.value(0)?;
162    parser.space();
163    if parser.at != parser.text.len() {
164        return Err(format!("unexpected text at byte {}", parser.at));
165    }
166    Ok(value)
167}
168
169struct Parser<'a> {
170    text: &'a [u8],
171    at: usize,
172}
173
174impl Parser<'_> {
175    fn space(&mut self) {
176        while self
177            .text
178            .get(self.at)
179            .is_some_and(|b| b.is_ascii_whitespace())
180        {
181            self.at += 1;
182        }
183    }
184
185    fn peek(&mut self) -> Option<u8> {
186        self.space();
187        self.text.get(self.at).copied()
188    }
189
190    fn value(&mut self, depth: usize) -> std::result::Result<Literal, String> {
191        if depth > MAX_DEPTH {
192            return Err("nested too deep".into());
193        }
194        match self.peek() {
195            Some(b'\'' | b'"') => self.string().map(Literal::Str),
196            Some(b'[') => self
197                .sequence(b']', depth)
198                .map(|(items, _)| Literal::List(items)),
199            Some(b'(') => {
200                let (items, comma) = self.sequence(b')', depth)?;
201                // `(x)` is x; `(x,)` is a tuple of one.
202                Ok(match (items.len(), comma) {
203                    (1, false) => items.into_iter().next().expect("one item"),
204                    _ => Literal::Tuple(items),
205                })
206            }
207            Some(b'{') => self.dict(depth),
208            Some(b'-' | b'+' | b'0'..=b'9') => self.int(),
209            Some(_) => self.word(),
210            None => Err("the header ends early".into()),
211        }
212    }
213
214    fn string(&mut self) -> std::result::Result<String, String> {
215        let quote = self.text[self.at];
216        self.at += 1;
217        let mut out = Vec::new();
218        loop {
219            let Some(&b) = self.text.get(self.at) else {
220                return Err("a string does not end".into());
221            };
222            self.at += 1;
223            match b {
224                b if b == quote => break,
225                b'\\' => {
226                    let Some(&e) = self.text.get(self.at) else {
227                        return Err("a string does not end".into());
228                    };
229                    self.at += 1;
230                    match e {
231                        b'n' => out.push(b'\n'),
232                        b't' => out.push(b'\t'),
233                        b'r' => out.push(b'\r'),
234                        b'0' => out.push(0),
235                        b'x' => {
236                            let hex = self
237                                .text
238                                .get(self.at..self.at + 2)
239                                .and_then(|h| std::str::from_utf8(h).ok())
240                                .and_then(|h| u8::from_str_radix(h, 16).ok())
241                                .ok_or("a bad \\x escape")?;
242                            self.at += 2;
243                            out.push(hex);
244                        }
245                        other => out.push(other),
246                    }
247                }
248                b => out.push(b),
249            }
250        }
251        Ok(String::from_utf8_lossy(&out).into_owned())
252    }
253
254    fn sequence(
255        &mut self,
256        close: u8,
257        depth: usize,
258    ) -> std::result::Result<(Vec<Literal>, bool), String> {
259        self.at += 1;
260        let mut items = Vec::new();
261        let mut comma = false;
262        loop {
263            if self.peek() == Some(close) {
264                self.at += 1;
265                return Ok((items, comma));
266            }
267            items.push(self.value(depth + 1)?);
268            comma = false;
269            match self.peek() {
270                Some(b',') => {
271                    self.at += 1;
272                    comma = true;
273                }
274                Some(b) if b == close => {}
275                _ => {
276                    return Err(format!(
277                        "expected , or {} at byte {}",
278                        close as char, self.at
279                    ));
280                }
281            }
282        }
283    }
284
285    fn dict(&mut self, depth: usize) -> std::result::Result<Literal, String> {
286        self.at += 1;
287        let mut pairs = Vec::new();
288        loop {
289            if self.peek() == Some(b'}') {
290                self.at += 1;
291                return Ok(Literal::Dict(pairs));
292            }
293            let key = self.value(depth + 1)?;
294            if self.peek() != Some(b':') {
295                return Err(format!("expected : at byte {}", self.at));
296            }
297            self.at += 1;
298            let value = self.value(depth + 1)?;
299            pairs.push((key, value));
300            match self.peek() {
301                Some(b',') => self.at += 1,
302                Some(b'}') => {}
303                _ => return Err(format!("expected , or }} at byte {}", self.at)),
304            }
305        }
306    }
307
308    fn int(&mut self) -> std::result::Result<Literal, String> {
309        let start = self.at;
310        if matches!(self.text.get(self.at), Some(b'-' | b'+')) {
311            self.at += 1;
312        }
313        while self.text.get(self.at).is_some_and(u8::is_ascii_digit) {
314            self.at += 1;
315        }
316        let digits = std::str::from_utf8(&self.text[start..self.at]).unwrap_or_default();
317        // Python 2 wrote long integers with an `L`: `(3L, 4L)`.
318        if matches!(self.text.get(self.at), Some(b'L' | b'l')) {
319            self.at += 1;
320        }
321        digits
322            .parse::<i128>()
323            .map(Literal::Int)
324            .map_err(|_| format!("a bad number at byte {start}"))
325    }
326
327    fn word(&mut self) -> std::result::Result<Literal, String> {
328        let start = self.at;
329        while self.text.get(self.at).is_some_and(u8::is_ascii_alphabetic) {
330            self.at += 1;
331        }
332        match &self.text[start..self.at] {
333            b"True" => Ok(Literal::Bool(true)),
334            b"False" => Ok(Literal::Bool(false)),
335            b"None" => Ok(Literal::None),
336            _ => Err(format!("unexpected text at byte {start}")),
337        }
338    }
339}
340
341// --- Element types -----------------------------------------------------------------
342
343/// An element type, as `descr` gives it.
344#[derive(Debug, Clone, PartialEq)]
345pub enum Dtype {
346    Simple(Simple),
347    /// Fields at offsets in an element of `itemsize` bytes.
348    Struct {
349        fields: Vec<FieldDef>,
350        itemsize: usize,
351    },
352}
353
354/// One field of a structured type.
355#[derive(Debug, Clone, PartialEq)]
356pub struct FieldDef {
357    pub name: String,
358    pub dtype: Dtype,
359    /// A subarray field's shape; empty for one value.
360    pub shape: Vec<usize>,
361    pub offset: usize,
362}
363
364/// A type with no fields: `<f8`, `|S10`, `<M8[ns]`.
365#[derive(Debug, Clone, PartialEq)]
366pub struct Simple {
367    /// The type string as written.
368    pub text: String,
369    pub kind: u8,
370    /// Bytes an element takes.
371    pub size: usize,
372    pub big: bool,
373    /// A datetime's or timedelta's unit and how many of it one step is.
374    pub unit: Option<(String, i64)>,
375}
376
377impl Dtype {
378    /// Bytes one element takes.
379    pub fn itemsize(&self) -> usize {
380        match self {
381            Self::Simple(s) => s.size,
382            Self::Struct { itemsize, .. } => *itemsize,
383        }
384    }
385
386    /// The type as NumPy writes it, for the Info panel.
387    pub fn describe(&self) -> String {
388        match self {
389            Self::Simple(s) => s.text.clone(),
390            Self::Struct { fields, itemsize } => {
391                format!("struct of {} fields, {itemsize} bytes", fields.len())
392            }
393        }
394    }
395}
396
397/// Parse a `descr`.
398pub fn parse_dtype(descr: &Literal) -> std::result::Result<Dtype, String> {
399    dtype_of(descr, 0)
400}
401
402fn dtype_of(descr: &Literal, depth: usize) -> std::result::Result<Dtype, String> {
403    if depth > MAX_DEPTH {
404        return Err("a structured type nested too deep".into());
405    }
406    match descr {
407        Literal::Str(text) => simple(text).map(Dtype::Simple),
408        Literal::List(items) => {
409            // Fields one after another; a field with no name is padding.
410            let mut fields = Vec::new();
411            let mut offset = 0usize;
412            for item in items {
413                let parts = item.items().ok_or("a field is not a tuple")?;
414                let (name, field_descr, shape) = match parts {
415                    [name, descr] => (name, descr, None),
416                    [name, descr, shape] => (name, descr, Some(shape)),
417                    _ => return Err("a field is not (name, type[, shape])".into()),
418                };
419                let name = field_name(name)?;
420                let dtype = dtype_of(field_descr, depth + 1)?;
421                let shape = shape.map(shape_of).transpose()?.unwrap_or_default();
422                let count = count_of(&shape)?;
423                let size = dtype
424                    .itemsize()
425                    .checked_mul(count)
426                    .ok_or("a field is too large")?;
427                if !name.is_empty() {
428                    fields.push(FieldDef {
429                        name,
430                        dtype,
431                        shape,
432                        offset,
433                    });
434                }
435                offset = offset.checked_add(size).ok_or("a type is too large")?;
436            }
437            Ok(Dtype::Struct {
438                fields,
439                itemsize: offset,
440            })
441        }
442        Literal::Dict(_) => {
443            // The dict form: names, formats, offsets and the element's size.
444            let list = |key: &str| {
445                descr
446                    .get(key)
447                    .and_then(Literal::items)
448                    .ok_or_else(|| format!("a structured type without {key}"))
449            };
450            let names = list("names")?;
451            let formats = list("formats")?;
452            if names.len() != formats.len() {
453                return Err("names and formats differ in length".into());
454            }
455            let offsets: Vec<usize> = match descr.get("offsets") {
456                Some(offsets) => offsets
457                    .items()
458                    .ok_or("offsets is not a list")?
459                    .iter()
460                    .map(|o| match o {
461                        Literal::Int(n) => usize::try_from(*n).map_err(|_| "a bad offset".into()),
462                        _ => Err("a bad offset".to_string()),
463                    })
464                    .collect::<std::result::Result<_, _>>()?,
465                None => Vec::new(),
466            };
467            let mut fields = Vec::new();
468            let mut next = 0usize;
469            for (i, (name, format)) in names.iter().zip(formats).enumerate() {
470                let name = field_name(name)?;
471                let (dtype, shape) = match format.items() {
472                    // `('<f8', (3,))`: a subarray in the dict form.
473                    Some([inner, shape]) => (dtype_of(inner, depth + 1)?, shape_of(shape)?),
474                    _ => (dtype_of(format, depth + 1)?, Vec::new()),
475                };
476                let size = dtype
477                    .itemsize()
478                    .checked_mul(count_of(&shape)?)
479                    .ok_or("a field is too large")?;
480                let offset = offsets.get(i).copied().unwrap_or(next);
481                next = offset.checked_add(size).ok_or("a type is too large")?;
482                if !name.is_empty() {
483                    fields.push(FieldDef {
484                        name,
485                        dtype,
486                        shape,
487                        offset,
488                    });
489                }
490            }
491            let end = fields
492                .iter()
493                .map(|f| f.offset + f.dtype.itemsize() * count_of(&f.shape).unwrap_or(1))
494                .max()
495                .unwrap_or(0);
496            let itemsize = match descr.get("itemsize") {
497                Some(Literal::Int(n)) => usize::try_from(*n).map_err(|_| "a bad itemsize")?,
498                _ => end,
499            };
500            if itemsize < end {
501                return Err(format!(
502                    "a field ends at byte {end}, past the element's {itemsize}"
503                ));
504            }
505            Ok(Dtype::Struct { fields, itemsize })
506        }
507        _ => Err("descr is not a type".into()),
508    }
509}
510
511/// A field's name: a string, or `(title, name)`.
512fn field_name(name: &Literal) -> std::result::Result<String, String> {
513    match name {
514        Literal::Str(s) => Ok(s.clone()),
515        Literal::Tuple(pair) if pair.len() == 2 => pair[1]
516            .as_str()
517            .map(str::to_string)
518            .ok_or_else(|| "a field's name is not text".into()),
519        _ => Err("a field's name is not text".into()),
520    }
521}
522
523fn shape_of(shape: &Literal) -> std::result::Result<Vec<usize>, String> {
524    match shape {
525        Literal::Int(n) => Ok(vec![usize::try_from(*n).map_err(|_| "a bad dimension")?]),
526        Literal::Tuple(items) | Literal::List(items) => items
527            .iter()
528            .map(|d| match d {
529                Literal::Int(n) => usize::try_from(*n).map_err(|_| "a bad dimension".into()),
530                _ => Err("a dimension is not a number".to_string()),
531            })
532            .collect(),
533        _ => Err("a shape is not a tuple".into()),
534    }
535}
536
537fn count_of(shape: &[usize]) -> std::result::Result<usize, String> {
538    shape
539        .iter()
540        .try_fold(1usize, |n, d| n.checked_mul(*d))
541        .ok_or_else(|| "a shape is too large".into())
542}
543
544/// Parse a type string such as `<f8`, `|S10`, `>i4` or `<M8[ns]`.
545fn simple(text: &str) -> std::result::Result<Simple, String> {
546    let bytes = text.as_bytes();
547    let (order, rest) = match bytes.first() {
548        Some(b @ (b'<' | b'>' | b'|' | b'=' | b'!')) => (*b, &text[1..]),
549        _ => (b'=', text),
550    };
551    let big = match order {
552        b'>' | b'!' => true,
553        b'=' => cfg!(target_endian = "big"),
554        _ => false,
555    };
556    let mut chars = rest.chars();
557    let kind = chars.next().ok_or("an empty type")?;
558    if !kind.is_ascii() {
559        return Err(format!("the type {text} is not one NumPy writes"));
560    }
561    let rest = chars.as_str();
562    let (digits, unit) = match rest.split_once('[') {
563        Some((digits, unit)) => (
564            digits,
565            Some(
566                unit.strip_suffix(']')
567                    .ok_or_else(|| format!("the type {text} has no ]"))?,
568            ),
569        ),
570        None => (rest, None),
571    };
572    let n: usize = if digits.is_empty() {
573        0
574    } else {
575        digits
576            .parse()
577            .map_err(|_| format!("the type {text} is not one NumPy writes"))?
578    };
579    let kind = kind as u8;
580    let size = match kind {
581        b'?' => 1,
582        b'U' => n.checked_mul(4).ok_or("a string type is too long")?,
583        b'O' => 8,
584        _ => n,
585    };
586    let unit = match unit {
587        Some(unit) => {
588            let split = unit
589                .find(|c: char| !c.is_ascii_digit())
590                .unwrap_or(unit.len());
591            let multiplier = match &unit[..split] {
592                "" => 1,
593                m => m.parse().map_err(|_| format!("the unit of {text}"))?,
594            };
595            Some((unit[split..].to_string(), multiplier))
596        }
597        None => None,
598    };
599    Ok(Simple {
600        text: text.to_string(),
601        kind,
602        size,
603        big,
604        unit,
605    })
606}
607
608/// How one simple type is read: what it is stored as and what it means, and how many
609/// values one element holds (two for a complex number).
610fn read_as(s: &Simple) -> std::result::Result<(Physical, usize, Logical, Option<Null>), String> {
611    let plain = Logical::Plain;
612    let unsupported = || format!("the type {} is not read", s.text);
613    Ok(match (s.kind, s.size) {
614        (b'b' | b'?', 1) => (Physical::Bool, 1, plain, None),
615        (b'i', 1 | 2 | 4 | 8) => (Physical::Signed(s.size as u8), 1, plain, None),
616        (b'u', 1 | 2 | 4 | 8) => (Physical::Unsigned(s.size as u8), 1, plain, None),
617        (b'f', 2 | 4 | 8) => (Physical::Float(s.size as u8), 1, plain, None),
618        // Real and imaginary parts, side by side.
619        (b'c', 8 | 16) => (Physical::Float((s.size / 2) as u8), 2, plain, None),
620        (b'S' | b'a', 1..) => (Physical::Text, 1, plain, None),
621        (b'U', 4..) => (Physical::Utf32 { big_endian: s.big }, 1, plain, None),
622        (b'V', 1..) => (Physical::Raw, 1, plain, None),
623        (b'O', _) => {
624            return Err(
625                "it holds Python objects (dtype O), which only unpickling could read; datui does not unpickle"
626                    .into(),
627            );
628        }
629        (b'M', 8) => {
630            let logical = match time_unit(s.unit.as_ref()) {
631                Some(Step::Days(1)) => Logical::Days { epoch_days: 0 },
632                Some(Step::Days(n)) => Logical::Timestamp {
633                    unit: TimeUnit::Milliseconds,
634                    multiplier: n.checked_mul(86_400_000).ok_or_else(unsupported)?,
635                    epoch: 0,
636                },
637                Some(Step::Unit(unit, multiplier)) => Logical::Timestamp {
638                    unit,
639                    multiplier,
640                    epoch: 0,
641                },
642                // Months and years have no fixed length, and finer than nanoseconds has
643                // no Polars type: their counts, as stored.
644                None => plain,
645            };
646            (Physical::Signed(8), 1, logical, Some(Null::Min))
647        }
648        (b'm', 8) => {
649            let logical = match time_unit(s.unit.as_ref()) {
650                Some(Step::Days(n)) => Logical::Duration {
651                    unit: TimeUnit::Milliseconds,
652                    multiplier: n.checked_mul(86_400_000).ok_or_else(unsupported)?,
653                },
654                Some(Step::Unit(unit, multiplier)) => Logical::Duration { unit, multiplier },
655                None => plain,
656            };
657            (Physical::Signed(8), 1, logical, Some(Null::Min))
658        }
659        _ => return Err(unsupported()),
660    })
661}
662
663enum Step {
664    /// A Polars unit, and how many of it one step is.
665    Unit(TimeUnit, i64),
666    /// Days, this many to a step.
667    Days(i64),
668}
669
670/// The step of a datetime or timedelta unit, when Polars has a type for it.
671fn time_unit(unit: Option<&(String, i64)>) -> Option<Step> {
672    let (name, n) = unit?;
673    let n = *n;
674    let ms = |per: i64| Some(Step::Unit(TimeUnit::Milliseconds, per.checked_mul(n)?));
675    match name.as_str() {
676        "ns" => Some(Step::Unit(TimeUnit::Nanoseconds, n)),
677        "us" | "\u{b5}s" => Some(Step::Unit(TimeUnit::Microseconds, n)),
678        "ms" => Some(Step::Unit(TimeUnit::Milliseconds, n)),
679        "s" => ms(1000),
680        "m" => ms(60_000),
681        "h" => ms(3_600_000),
682        "D" => Some(Step::Days(n)),
683        "W" => Some(Step::Days(n.checked_mul(7)?)),
684        _ => None,
685    }
686}
687
688// --- The header ---------------------------------------------------------------------
689
690/// What a `.npy` header says.
691#[derive(Debug, Clone, PartialEq)]
692pub struct Header {
693    pub version: (u8, u8),
694    pub dtype: Dtype,
695    pub fortran: bool,
696    pub shape: Vec<usize>,
697    /// Bytes from the start of the file to the first element.
698    pub data_offset: usize,
699}
700
701/// Bytes before the header's text, and the header's length, from the file's first 12
702/// bytes; `None` when they are not a `.npy` file's.
703pub fn header_len(head: &[u8]) -> std::result::Result<(usize, usize), String> {
704    if !looks_like(head) {
705        return Err("not a NumPy .npy file: no \\x93NUMPY at the start".into());
706    }
707    let major = *head.get(6).ok_or("the header is cut short")?;
708    let (prefix, len) = match major {
709        1 => (
710            10,
711            u16::from_le_bytes(
712                head.get(8..10)
713                    .ok_or("the header is cut short")?
714                    .try_into()
715                    .expect("two bytes"),
716            ) as usize,
717        ),
718        2 | 3 => (
719            12,
720            u32::from_le_bytes(
721                head.get(8..12)
722                    .ok_or("the header is cut short")?
723                    .try_into()
724                    .expect("four bytes"),
725            ) as usize,
726        ),
727        v => {
728            return Err(format!(
729                "format version {v} is not one datui reads (1, 2 or 3)"
730            ));
731        }
732    };
733    // NumPy's own reader refuses past 10,000 bytes unless told otherwise; a structured
734    // type of many fields can be longer, and the default is still small.
735    let most = crate::limits::get().npy_header_bytes.bytes();
736    if len as u64 > most {
737        let size = crate::numfmt::bytes;
738        return Err(format!(
739            "its header is {}; datui reads headers up to {}, and limits.npy_header_bytes raises it",
740            size(len as u64),
741            size(most)
742        ));
743    }
744    Ok((prefix, len))
745}
746
747/// Parse the header at the start of `bytes`, which must hold all of it.
748pub fn parse_header(bytes: &[u8]) -> std::result::Result<Header, String> {
749    let (prefix, len) = header_len(bytes)?;
750    let text = bytes
751        .get(prefix..prefix + len)
752        .ok_or("the header is cut short")?;
753    // Version 3 headers are UTF-8, earlier ones Latin-1; the keys and types are ASCII.
754    let text: String = match bytes[6] {
755        3 => String::from_utf8_lossy(text).into_owned(),
756        _ => text.iter().map(|&b| char::from(b)).collect(),
757    };
758    let dict = parse_literal(text.trim_end_matches(['\n', ' ', '\0']))
759        .map_err(|e| format!("its header is not a Python dict: {e}"))?;
760    if !matches!(dict, Literal::Dict(_)) {
761        return Err("its header is not a Python dict".into());
762    }
763    let descr = dict.get("descr").ok_or("its header has no descr")?;
764    let dtype = parse_dtype(descr)?;
765    let fortran = match dict.get("fortran_order") {
766        Some(Literal::Bool(b)) => *b,
767        None => false,
768        _ => return Err("fortran_order is not True or False".into()),
769    };
770    let shape = shape_of(dict.get("shape").ok_or("its header has no shape")?)?;
771    count_of(&shape)?;
772    Ok(Header {
773        version: (bytes[6], bytes[7]),
774        dtype,
775        fortran,
776        shape,
777        data_offset: prefix + len,
778    })
779}
780
781/// The shape as Python writes it: `(3,)`, `(2, 4)`, `()`.
782pub fn shape_text(shape: &[usize]) -> String {
783    match shape {
784        [one] => format!("({one},)"),
785        _ => format!(
786            "({})",
787            shape
788                .iter()
789                .map(usize::to_string)
790                .collect::<Vec<_>>()
791                .join(", ")
792        ),
793    }
794}
795
796// --- Columns ------------------------------------------------------------------------
797
798/// The columns an array is read as, their starts counted from its first element, and
799/// its rows.
800#[derive(Debug)]
801pub struct Plan {
802    pub columns: Vec<ColumnLayout>,
803    pub rows: usize,
804    /// Bytes the elements take, as the shape says.
805    pub bytes: u128,
806}
807
808/// The columns of an array of `header`, named `name` where it is one column.
809pub fn plan(header: &Header, name: &str) -> std::result::Result<Plan, String> {
810    let itemsize = header.dtype.itemsize();
811    if itemsize == 0 {
812        return Err("its elements take no bytes".into());
813    }
814    let elements = count_of(&header.shape)?;
815    let bytes = elements as u128 * itemsize as u128;
816    let refuse_shape = || {
817        format!(
818            "it has {} dimensions, shape {}; datui reads 1-D and 2-D arrays",
819            header.shape.len(),
820            shape_text(&header.shape)
821        )
822    };
823    match (&header.dtype, header.shape.as_slice()) {
824        (_, [] | [_]) => {
825            let mut columns = Vec::new();
826            // A structured array's fields are its columns, named as they are.
827            let name = match header.dtype {
828                Dtype::Struct { .. } => "",
829                Dtype::Simple(_) => name,
830            };
831            leaves(&header.dtype, name, 0, itemsize, &mut columns, 0)?;
832            Ok(Plan {
833                columns,
834                rows: header.shape.first().copied().unwrap_or(1),
835                bytes,
836            })
837        }
838        (Dtype::Simple(_), [rows, cols]) => {
839            let (rows, cols) = (*rows, *cols);
840            let mut columns = Vec::new();
841            if cols > MAX_GRID {
842                if header.fortran {
843                    return Err(format!(
844                        "it is {} in Fortran order, with more than {MAX_GRID} columns",
845                        shape_text(&header.shape)
846                    ));
847                }
848                let mut leaf = Vec::new();
849                leaves(&header.dtype, "values", 0, itemsize, &mut leaf, 0)?;
850                let mut column = leaf.remove(0);
851                column.count = column.count.checked_mul(cols).ok_or("a row is too large")?;
852                column.stride = itemsize.checked_mul(cols).ok_or("a row is too large")?;
853                columns.push(column);
854            } else {
855                for c in 0..cols {
856                    let mut leaf = Vec::new();
857                    leaves(&header.dtype, &c.to_string(), 0, itemsize, &mut leaf, 0)?;
858                    let mut column = leaf.remove(0);
859                    if header.fortran {
860                        // Column-major: a column's elements are side by side.
861                        column.start = c * rows * itemsize;
862                        column.stride = itemsize;
863                    } else {
864                        column.start = c * itemsize;
865                        column.stride = itemsize * cols;
866                    }
867                    columns.push(column);
868                }
869            }
870            Ok(Plan {
871                columns,
872                rows,
873                bytes,
874            })
875        }
876        (Dtype::Struct { .. }, [_, _]) => Err(format!(
877            "it is a structured array of shape {}; datui reads structured arrays of one dimension",
878            shape_text(&header.shape)
879        )),
880        _ => Err(refuse_shape()),
881    }
882}
883
884/// The columns of an element of `dtype` at `offset` into a record of `stride` bytes,
885/// named under `name`.
886fn leaves(
887    dtype: &Dtype,
888    name: &str,
889    offset: usize,
890    stride: usize,
891    out: &mut Vec<ColumnLayout>,
892    depth: usize,
893) -> std::result::Result<(), String> {
894    if depth > MAX_DEPTH {
895        return Err("a structured type nested too deep".into());
896    }
897    match dtype {
898        Dtype::Simple(s) => {
899            let (physical, values, logical, null) =
900                read_as(s).map_err(|e| format!("column \"{name}\": {e}"))?;
901            let width = s.size / values;
902            out.push(ColumnLayout {
903                name: name.into(),
904                source: 0,
905                start: offset,
906                stride,
907                width,
908                count: values,
909                physical,
910                big_endian: s.big,
911                null,
912                logical,
913            });
914        }
915        Dtype::Struct { fields, .. } => {
916            for field in fields {
917                let full = match name {
918                    "" => field.name.clone(),
919                    _ => format!("{name}.{}", field.name),
920                };
921                let count = count_of(&field.shape)?;
922                let at = offset
923                    .checked_add(field.offset)
924                    .ok_or("a field is too far in")?;
925                match &field.dtype {
926                    Dtype::Simple(_) => {
927                        leaves(&field.dtype, &full, at, stride, out, depth + 1)?;
928                        if count != 1
929                            && let Some(column) = out.last_mut()
930                        {
931                            // A subarray field: its values in one Array cell.
932                            column.count = column
933                                .count
934                                .checked_mul(count)
935                                .ok_or("a subarray is too large")?;
936                        }
937                    }
938                    Dtype::Struct { .. } if count == 1 => {
939                        leaves(&field.dtype, &full, at, stride, out, depth + 1)?;
940                    }
941                    Dtype::Struct { itemsize, .. } => {
942                        for k in 0..count {
943                            let at = k
944                                .checked_mul(*itemsize)
945                                .and_then(|o| at.checked_add(o))
946                                .ok_or("a field is too far in")?;
947                            leaves(
948                                &field.dtype,
949                                &format!("{full}[{k}]"),
950                                at,
951                                stride,
952                                out,
953                                depth + 1,
954                            )?;
955                            if out.len() > MAX_COLUMNS {
956                                break;
957                            }
958                        }
959                    }
960                }
961                if out.len() > MAX_COLUMNS {
962                    return Err(format!("its fields make more than {MAX_COLUMNS} columns"));
963                }
964            }
965        }
966    }
967    Ok(())
968}
969
970// --- Reading ------------------------------------------------------------------------
971
972/// An array read in place: its records and what to say about it.
973pub struct Array {
974    pub records: Arc<FixedRecords>,
975    pub header: Header,
976    pub notes: Vec<String>,
977}
978
979/// The array whose `.npy` bytes start `at` bytes into `bytes` (the whole of a file, or
980/// a member stored in an archive) and run for `len` bytes, its one column named `name`.
981pub fn open_in(
982    bytes: Arc<Bytes>,
983    at: usize,
984    len: usize,
985    name: &str,
986) -> std::result::Result<Array, String> {
987    let end = at.checked_add(len).filter(|&end| end <= bytes.len());
988    let Some(end) = end else {
989        return Err("the array runs past the end of the file".into());
990    };
991    let npy = &bytes.as_slice()[at..end];
992    let (prefix, header_bytes) = header_len(npy)?;
993    if npy.len() < prefix + header_bytes {
994        return Err("the header is cut short".into());
995    }
996    let header = parse_header(npy)?;
997    let plan = plan(&header, name)?;
998    let data = header.data_offset;
999    let on_hand = (npy.len() - data) as u128;
1000    let mut notes = Vec::new();
1001    if on_hand < plan.bytes {
1002        notes.push(format!(
1003            "file cut short: {} of {} bytes for shape {} {} rows past the end left out",
1004            crate::numfmt::group_chrome(on_hand as usize),
1005            crate::numfmt::group_chrome(usize::try_from(plan.bytes).unwrap_or(usize::MAX)),
1006            shape_text(&header.shape),
1007            crate::glyphs::get().middot
1008        ));
1009    }
1010    if plan.rows > crate::formats::row_index::MAX_ROWS {
1011        notes.push(format!(
1012            "{} rows {} first {} shown",
1013            crate::numfmt::group_chrome(plan.rows),
1014            crate::glyphs::get().middot,
1015            crate::numfmt::group_chrome(crate::formats::row_index::MAX_ROWS)
1016        ));
1017    }
1018    // A cell past what is on hand would be read past the array's end in an archive,
1019    // into the next member: the rows are cut to the array's own bytes.
1020    let rows_on_hand = |c: &ColumnLayout| -> usize {
1021        let cell = c.width * c.count;
1022        let room = (end - at).saturating_sub(data + c.start);
1023        if room < cell || c.stride == 0 {
1024            0
1025        } else {
1026            (room - cell) / c.stride + 1
1027        }
1028    };
1029    let rows = plan
1030        .columns
1031        .iter()
1032        .map(rows_on_hand)
1033        .min()
1034        .unwrap_or(plan.rows)
1035        .min(plan.rows);
1036    let columns = plan
1037        .columns
1038        .into_iter()
1039        .map(|c| ColumnLayout {
1040            start: c.start + at + data,
1041            ..c
1042        })
1043        .collect();
1044    let records = FixedRecords::new(vec![bytes], columns, rows).map_err(|e| e.to_string())?;
1045    Ok(Array {
1046        records: Arc::new(records),
1047        header,
1048        notes,
1049    })
1050}
1051
1052/// The array of the `.npy` file at `path`.
1053pub fn open_file(path: &Path) -> Result<Array> {
1054    let bytes = Bytes::map(path).map_err(|e| in_file(path, e.into()))?;
1055    let len = bytes.len();
1056    open_in(Arc::new(bytes), 0, len, &stem(path)).map_err(|e| FileError::new(path, e).into())
1057}
1058
1059fn stem(path: &Path) -> String {
1060    path.file_stem()
1061        .map(|s| s.to_string_lossy().into_owned())
1062        .unwrap_or_else(|| "values".to_string())
1063}
1064
1065/// What the Info panel's NumPy tab says of an array.
1066pub fn detail(header: &Header, archive: Option<(&Path, usize)>, compressed: bool) -> Detail {
1067    let mut lines = vec![
1068        format!("Shape: {}", shape_text(&header.shape)),
1069        format!("Type: {}", header.dtype.describe()),
1070        format!(
1071            "Order: {}",
1072            if header.fortran {
1073                "Fortran (column-major)"
1074            } else {
1075                "C (row-major)"
1076            }
1077        ),
1078        format!("Version: {}.{}", header.version.0, header.version.1),
1079    ];
1080    if let Some((archive, arrays)) = archive {
1081        lines.push(format!(
1082            "Archive: {}, {}{}",
1083            archive.display(),
1084            crate::formats::text_formats::count(arrays as u64, "array", "arrays"),
1085            if compressed {
1086                "; this one compressed, read from a copy"
1087            } else {
1088                ""
1089            }
1090        ));
1091    }
1092    let mut list = Vec::new();
1093    if let Dtype::Struct { fields, .. } = &header.dtype {
1094        let rows = fields.iter().map(|f| {
1095            let mut text = f.dtype.describe();
1096            if !f.shape.is_empty() {
1097                text.push_str(&format!(" {}", shape_text(&f.shape)));
1098            }
1099            text.push_str(&format!(" at byte {}", f.offset));
1100            (f.name.clone(), MetaValue::Text(text))
1101        });
1102        list = crate::formats::text_formats::capped_list(rows, fields.len());
1103    }
1104    Detail {
1105        tab: crate::formats::text_formats::tab(crate::FileFormat::Numpy),
1106        lines,
1107        list_title: "Fields",
1108        list,
1109        first: false,
1110        ..Default::default()
1111    }
1112}
1113
1114// --- Archives -----------------------------------------------------------------------
1115
1116/// One array of an archive.
1117#[derive(Debug, Clone)]
1118pub struct Member {
1119    /// The name `np.load` gives it: the file name in the archive, less `.npy`.
1120    pub name: String,
1121    /// The entry's index in the archive.
1122    pub index: usize,
1123    /// Stored as it is, at this offset in the archive; `None` when compressed.
1124    pub stored_at: Option<u64>,
1125    pub size: u64,
1126}
1127
1128fn archive_of(path: &Path) -> Result<::zip::ZipArchive<std::fs::File>> {
1129    let file = std::fs::File::open(path).map_err(|e| in_file(path, e.into()))?;
1130    ::zip::ZipArchive::new(file)
1131        .map_err(|e| FileError::new(path, format!("not a NumPy archive: {e}")).into())
1132}
1133
1134/// The arrays of the archive at `path`, in the archive's order.
1135pub fn members(path: &Path) -> Result<Vec<Member>> {
1136    let mut archive = archive_of(path)?;
1137    let mut members = Vec::new();
1138    for index in 0..archive.len().min(MAX_MEMBERS) {
1139        let Ok(entry) = archive.by_index_raw(index) else {
1140            continue;
1141        };
1142        let Some(name) = entry.name().strip_suffix(".npy") else {
1143            continue;
1144        };
1145        if entry.is_dir() || name.is_empty() {
1146            continue;
1147        }
1148        let stored = entry.compression() == ::zip::CompressionMethod::Stored;
1149        members.push(Member {
1150            name: name.to_string(),
1151            index,
1152            stored_at: if stored { entry.data_start() } else { None },
1153            size: entry.size(),
1154        });
1155    }
1156    Ok(members)
1157}
1158
1159/// The header of `member`, read from its first bytes, decompressed when it is
1160/// compressed: never more than the header.
1161fn member_header(path: &Path, member: &Member) -> std::result::Result<Header, String> {
1162    let mut archive = archive_of(path).map_err(|e| e.to_string())?;
1163    let entry = archive.by_index(member.index).map_err(|e| e.to_string())?;
1164    let mut reader = entry.take(12);
1165    let mut head = Vec::with_capacity(12);
1166    reader.read_to_end(&mut head).map_err(|e| e.to_string())?;
1167    let (prefix, len) = header_len(&head)?;
1168    let mut entry = reader.into_inner();
1169    let mut rest = Vec::with_capacity(len);
1170    (&mut entry)
1171        .take((prefix + len - head.len()) as u64)
1172        .read_to_end(&mut rest)
1173        .map_err(|e| e.to_string())?;
1174    head.extend(rest);
1175    parse_header(&head)
1176}
1177
1178/// The arrays of an archive as tables, with their columns, for the home screen.
1179pub fn tables(path: &Path) -> Result<Vec<Table>> {
1180    let members = members(path)?;
1181    Ok(members
1182        .iter()
1183        .map(|member| {
1184            let columns: Vec<String> = member_header(path, member)
1185                .ok()
1186                .and_then(|header| plan(&header, &member.name).ok())
1187                .map(|plan| plan.columns.iter().map(|c| c.name.to_string()).collect())
1188                .unwrap_or_default();
1189            Table::plain(&member.name, "array", columns)
1190        })
1191        .collect())
1192}
1193
1194/// The columns an array would open with, from its header: a `.npy` file, or an
1195/// archive's `member`. `None` for one that would not open.
1196pub fn schema_preview(
1197    path: &Path,
1198    member: Option<&str>,
1199) -> Option<crate::home::discover::SchemaPreview> {
1200    let (header, name) = match member {
1201        Some(name) => {
1202            let found = members(path).ok()?.into_iter().find(|m| m.name == name)?;
1203            (member_header(path, &found).ok()?, found.name)
1204        }
1205        None if is_archive_name(path) => {
1206            let all = members(path).ok()?;
1207            let [one] = all.as_slice() else {
1208                return None;
1209            };
1210            (member_header(path, one).ok()?, one.name.clone())
1211        }
1212        None => {
1213            let mut file = std::fs::File::open(path).ok()?;
1214            let mut head = [0u8; 12];
1215            file.read_exact(&mut head).ok()?;
1216            let (prefix, len) = header_len(&head).ok()?;
1217            let mut bytes = head.to_vec();
1218            file.take((prefix + len - 12) as u64)
1219                .read_to_end(&mut bytes)
1220                .ok()?;
1221            (parse_header(&bytes).ok()?, stem(path))
1222        }
1223    };
1224    let plan = plan(&header, &name).ok()?;
1225    Some(
1226        plan.columns
1227            .iter()
1228            .map(|c| (c.name.to_string(), c.dtype()))
1229            .collect(),
1230    )
1231}
1232
1233/// What the open carries to the dataset for `array`.
1234fn opened(
1235    array: Array,
1236    archive: Option<(&Path, usize)>,
1237    compressed: bool,
1238) -> crate::formats::members::Opened {
1239    let rows = array.records.rows();
1240    crate::formats::members::Opened {
1241        window: Some((array.records.clone(), rows)),
1242        detail: Some(Arc::new(detail(&array.header, archive, compressed))),
1243        other_tables: Vec::new(),
1244        notes: array
1245            .notes
1246            .into_iter()
1247            .map(|summary| crate::formats::text_formats::note(summary, "the array".to_string()))
1248            .collect(),
1249        units: Vec::new(),
1250        indexing: None,
1251        numbering: None,
1252    }
1253}
1254
1255/// Decompress the archive's array `name` into a temporary `.npy` file written through
1256/// `writer`, counting the archive's bytes read in `read`, and open it.
1257pub(crate) fn convert(
1258    path: &Path,
1259    name: &str,
1260    options: &crate::OpenOptions,
1261    writer: &crate::loading::unfinished::Writer,
1262    read: &std::sync::atomic::AtomicU64,
1263) -> Result<(
1264    crate::cloud::download::TempDownload,
1265    LazyFrame,
1266    crate::formats::members::Opened,
1267)> {
1268    use std::io::Write;
1269    use std::sync::atomic::Ordering;
1270    let all = members(path)?;
1271    let member = all
1272        .iter()
1273        .find(|m| m.name == name)
1274        .ok_or_else(|| FileError::new(path, format!("no array \"{name}\"")))?;
1275    let mut archive = archive_of(path)?;
1276    let entry = archive
1277        .by_index(member.index)
1278        .map_err(|e| FileError::new(path, format!("array \"{name}\": {e}")))?;
1279    let compressed = entry.compressed_size();
1280    let Some((mut file, claim)) = writer.create(|| {
1281        crate::cloud::download::TempDownload::create(options.temp_dir.as_deref(), Some("npy"))
1282    })?
1283    else {
1284        return Err(eyre!("Reading was stopped."));
1285    };
1286    // No more than the archive says the array is, so a member that inflates past its
1287    // stated size stops there rather than filling the disk.
1288    let mut reader = entry.take(member.size);
1289    let mut chunk = vec![0u8; 1 << 16];
1290    let mut written = 0u64;
1291    loop {
1292        if writer.stopped() {
1293            return Err(eyre!("Reading was stopped."));
1294        }
1295        let n = match reader.read(&mut chunk) {
1296            Ok(0) => break,
1297            Ok(n) => n,
1298            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
1299            Err(e) => return Err(FileError::new(path, format!("array \"{name}\": {e}")).into()),
1300        };
1301        file.write_all(&chunk[..n])?;
1302        written += n as u64;
1303        if member.size > 0 {
1304            // The compressed bytes read so far, in proportion.
1305            let share = (compressed as u128 * n as u128 / member.size as u128) as u64;
1306            read.fetch_add(share, Ordering::Relaxed);
1307        }
1308    }
1309    file.flush()?;
1310    let held = crate::cloud::download::TempDownload::held(file, Some(claim));
1311    let bytes = Bytes::map(held.path())?;
1312    let array = open_in(Arc::new(bytes), 0, written as usize, name)
1313        .map_err(|e| FileError::new(path, format!("array \"{name}\" is not read: {e}")))?;
1314    let tables = arrays(&all);
1315    let lf = array.records.lazy();
1316    let mut opened = opened(array, Some((path, all.len())), true);
1317    opened.other_tables = crate::formats::members::others(&tables, name);
1318    Ok((held, lf, opened))
1319}
1320
1321/// The scan of a NumPy file: an `.npy` file's array, or the array of an archive
1322/// `--table` names, or its only one, read in place; a compressed one decompressed
1323/// first; or none yet when the archive has several.
1324fn scan(input: crate::formats::readers::ScanIn<'_>) -> Result<crate::loading::scan::Scan> {
1325    let path = &input.path().to_path_buf();
1326    let wanted = input.options.table.as_deref();
1327    let mut head = [0u8; 8];
1328    let is_npy = std::fs::File::open(path)
1329        .and_then(|mut f| f.read_exact(&mut head))
1330        .is_ok()
1331        && looks_like(&head);
1332    if is_npy {
1333        if let Some(wanted) = wanted {
1334            return Err(FileError::new(
1335                path,
1336                format!(
1337                    "the file is one array; --table \"{wanted}\" picks an array of an .npz archive."
1338                ),
1339            )
1340            .into());
1341        }
1342        let array = open_file(path)?;
1343        let lf = array.records.lazy();
1344        return Ok(opened(array, None, false).scan(input, lf));
1345    }
1346    if !is_archive_name(path) {
1347        return Err(FileError::new(
1348            path,
1349            "not a NumPy file: an .npy file starts with \\x93NUMPY, and an .npz is a zip of them.",
1350        )
1351        .into());
1352    }
1353    let members = members(path)?;
1354    let tables = arrays(&members);
1355    let picked = match crate::formats::members::pick(tables.clone(), wanted, path, "")? {
1356        crate::formats::sqlite::Pick::One(table) => table,
1357        crate::formats::sqlite::Pick::Several(tables) => {
1358            return Ok(crate::formats::members::several(&input, tables));
1359        }
1360    };
1361    let member = members
1362        .iter()
1363        .find(|m| m.name == picked.name)
1364        .expect("picked from the members");
1365    let Some(at) = member.stored_at else {
1366        return Ok(crate::loading::scan::Scan::Unpack {
1367            file: path.clone(),
1368            member: member.name.clone(),
1369            format: input.format,
1370        });
1371    };
1372    let bytes = Bytes::map(path).map_err(|e| in_file(path, e.into()))?;
1373    let at = usize::try_from(at).map_err(|_| FileError::new(path, "the file is too large"))?;
1374    let len = usize::try_from(member.size).unwrap_or(usize::MAX);
1375    let array = open_in(Arc::new(bytes), at, len, &member.name)
1376        .map_err(|e| FileError::new(path, format!("array \"{}\" is not read: {e}", member.name)))?;
1377    let lf = array.records.lazy();
1378    let mut opened = opened(array, Some((path, members.len())), false);
1379    opened.other_tables = crate::formats::members::others(&tables, &member.name);
1380    Ok(opened.scan(input, lf))
1381}
1382
1383/// An archive's arrays as tables, without their columns.
1384fn arrays(members: &[Member]) -> Vec<Table> {
1385    members
1386        .iter()
1387        .map(|m| Table::plain(&m.name, "array", Vec::<String>::new()))
1388        .collect()
1389}
1390
1391#[cfg(test)]
1392mod tests {
1393    use super::*;
1394
1395    /// A header past `limits.npy_header_bytes` is refused, naming the key.
1396    #[test]
1397    fn a_header_past_the_limit_names_its_key() {
1398        let mut head = b"\x93NUMPY\x02\x00".to_vec();
1399        let past = crate::limits::get().npy_header_bytes.bytes() as u32 + 1;
1400        head.extend_from_slice(&past.to_le_bytes());
1401        let refused = header_len(&head).unwrap_err();
1402        assert!(refused.contains("limits.npy_header_bytes"), "{refused}");
1403        head[8..12].copy_from_slice(&64u32.to_le_bytes());
1404        assert_eq!(header_len(&head), Ok((12, 64)));
1405    }
1406
1407    /// Every way an array is refused names the file, in the one shape.
1408    #[test]
1409    fn errors_name_the_file() {
1410        use crate::formats::readers::bad_input::{assert_shape, each_names_its_file, opening};
1411        let mut v9 = npy("'<f8'", false, "(1,)", &[0; 8]);
1412        v9[6] = 9;
1413        each_names_its_file(
1414            crate::FileFormat::Numpy,
1415            &[
1416                ("text.npy", b"hello there", "Not a NumPy file"),
1417                ("text.npz", b"hello there", "Not a NumPy archive"),
1418                ("v9.npy", &v9, "version 9"),
1419                ("type.npy", &npy("'<q9'", false, "(1,)", &[0; 8]), "<q9"),
1420            ],
1421        );
1422        let dir = tempfile::tempdir().unwrap();
1423        let options = crate::OpenOptions {
1424            table: Some("x".into()),
1425            ..Default::default()
1426        };
1427        let one = npy("'<f8'", false, "(1,)", &[0; 8]);
1428        let message = opening(
1429            dir.path(),
1430            "a.npy",
1431            &one,
1432            crate::FileFormat::Numpy,
1433            &options,
1434        )
1435        .expect("--table refused");
1436        eprintln!("{message}");
1437        assert_shape(&message, &dir.path().join("a.npy"));
1438        assert!(message.contains("--table \"x\""), "{message}");
1439    }
1440
1441    fn npy(descr: &str, fortran: bool, shape: &str, data: &[u8]) -> Vec<u8> {
1442        let mut header = format!(
1443            "{{'descr': {descr}, 'fortran_order': {}, 'shape': {shape}, }}",
1444            if fortran { "True" } else { "False" }
1445        );
1446        while (10 + header.len() + 1) % 64 != 0 {
1447            header.push(' ');
1448        }
1449        header.push('\n');
1450        let mut out = MAGIC.to_vec();
1451        out.extend([1, 0]);
1452        out.extend((header.len() as u16).to_le_bytes());
1453        out.extend(header.as_bytes());
1454        out.extend(data);
1455        out
1456    }
1457
1458    fn read(bytes: Vec<u8>, name: &str) -> DataFrame {
1459        let len = bytes.len();
1460        let array = open_in(Arc::new(Bytes::Owned(bytes)), 0, len, name).unwrap();
1461        array.records.lazy().collect().unwrap()
1462    }
1463
1464    #[test]
1465    fn literals_parse() {
1466        let parsed = parse_literal("{'a': [('x', '<f8', (3,)), (('t', 'y'), '|b1')], 'b': (2L,), 'c': True, 'd': None, 'e': -4}")
1467            .unwrap();
1468        assert_eq!(
1469            parsed.get("b"),
1470            Some(&Literal::Tuple(vec![Literal::Int(2)]))
1471        );
1472        assert_eq!(parsed.get("c"), Some(&Literal::Bool(true)));
1473        assert_eq!(parsed.get("e"), Some(&Literal::Int(-4)));
1474        assert_eq!(parse_literal("(1)").unwrap(), Literal::Int(1));
1475        assert_eq!(parse_literal("()").unwrap(), Literal::Tuple(Vec::new()));
1476        assert!(parse_literal("{'a': 1").is_err());
1477        assert!(parse_literal(&"[".repeat(100)).is_err());
1478        assert!(parse_literal("__import__('os')").is_err());
1479    }
1480
1481    #[test]
1482    fn a_vector_is_one_column_named_for_the_array() {
1483        let data: Vec<u8> = [1.5f64, -2.0, 3.25]
1484            .iter()
1485            .flat_map(|v| v.to_le_bytes())
1486            .collect();
1487        let df = read(npy("'<f8'", false, "(3,)", &data), "prices");
1488        assert_eq!(df.get_column_names(), ["prices"]);
1489        assert_eq!(
1490            df.column("prices").unwrap().f64().unwrap().to_vec(),
1491            [Some(1.5), Some(-2.0), Some(3.25)]
1492        );
1493    }
1494
1495    #[test]
1496    fn a_grid_in_either_order() {
1497        // [[1, 2, 3], [4, 5, 6]] as int32.
1498        let c: Vec<u8> = [1i32, 2, 3, 4, 5, 6]
1499            .iter()
1500            .flat_map(|v| v.to_le_bytes())
1501            .collect();
1502        let f: Vec<u8> = [1i32, 4, 2, 5, 3, 6]
1503            .iter()
1504            .flat_map(|v| v.to_le_bytes())
1505            .collect();
1506        for df in [
1507            read(npy("'<i4'", false, "(2, 3)", &c), "g"),
1508            read(npy("'<i4'", true, "(2, 3)", &f), "g"),
1509        ] {
1510            assert_eq!(df.get_column_names(), ["0", "1", "2"]);
1511            assert_eq!(
1512                df.column("2").unwrap().i32().unwrap().to_vec(),
1513                [Some(3), Some(6)]
1514            );
1515        }
1516    }
1517
1518    #[test]
1519    fn structured_fields_mixed_endian_subarrays_and_padding() {
1520        // ts <u8, px >f4 (big-endian), 3 bytes of padding, v <i2 (2,)
1521        let mut data = Vec::new();
1522        for (ts, px, v) in [(1u64, 0.5f32, [7i16, 8]), (2, 1.5, [9, 10])] {
1523            data.extend(ts.to_le_bytes());
1524            data.extend(px.to_be_bytes());
1525            data.extend([0u8; 3]);
1526            data.extend(v.iter().flat_map(|x| x.to_le_bytes()));
1527        }
1528        let descr = "[('ts', '<u8'), ('px', '>f4'), ('', '|V3'), ('v', '<i2', (2,))]";
1529        let df = read(npy(descr, false, "(2,)", &data), "t");
1530        assert_eq!(df.get_column_names(), ["ts", "px", "v"]);
1531        assert_eq!(
1532            df.column("px").unwrap().f32().unwrap().to_vec(),
1533            [Some(0.5), Some(1.5)]
1534        );
1535        assert_eq!(
1536            df.column("v").unwrap().dtype(),
1537            &DataType::Array(Box::new(DataType::Int16), 2)
1538        );
1539    }
1540
1541    #[test]
1542    fn the_dict_form_honors_offsets_and_itemsize() {
1543        // a at 4, b at 0, in 12-byte elements.
1544        let mut data = Vec::new();
1545        for (a, b) in [(10i32, 20i32), (30, 40)] {
1546            data.extend(b.to_le_bytes());
1547            data.extend(a.to_le_bytes());
1548            data.extend([0u8; 4]);
1549        }
1550        let descr =
1551            "{'names': ['a', 'b'], 'formats': ['<i4', '<i4'], 'offsets': [4, 0], 'itemsize': 12}";
1552        let df = read(npy(descr, false, "(2,)", &data), "t");
1553        assert_eq!(
1554            df.column("a").unwrap().i32().unwrap().to_vec(),
1555            [Some(10), Some(30)]
1556        );
1557        assert_eq!(
1558            df.column("b").unwrap().i32().unwrap().to_vec(),
1559            [Some(20), Some(40)]
1560        );
1561    }
1562
1563    #[test]
1564    fn strings_times_and_bools() {
1565        let mut data = Vec::new();
1566        // U3 "hé", S4 "ab", M8[s] 86400 (1970-01-02), m8[ms] 1500, b1 true.
1567        for c in "hé".chars().chain(std::iter::once('\0')) {
1568            data.extend((c as u32).to_le_bytes());
1569        }
1570        data.extend(b"ab\0\0");
1571        data.extend(86_400i64.to_le_bytes());
1572        data.extend(1500i64.to_le_bytes());
1573        data.push(1);
1574        let descr = "[('u', '<U3'), ('s', '|S4'), ('t', '<M8[s]'), ('d', '<m8[ms]'), ('b', '|b1')]";
1575        let df = read(npy(descr, false, "(1,)", &data), "t");
1576        assert_eq!(df.column("u").unwrap().str().unwrap().get(0), Some("hé"));
1577        assert_eq!(df.column("s").unwrap().str().unwrap().get(0), Some("ab"));
1578        assert_eq!(
1579            df.column("t").unwrap().dtype(),
1580            &DataType::Datetime(TimeUnit::Milliseconds, None)
1581        );
1582        assert_eq!(
1583            df.column("d").unwrap().dtype(),
1584            &DataType::Duration(TimeUnit::Milliseconds)
1585        );
1586        assert_eq!(df.column("b").unwrap().bool().unwrap().get(0), Some(true));
1587    }
1588
1589    #[test]
1590    fn what_is_refused() {
1591        let refused = |descr: &str, shape: &str| {
1592            let bytes = npy(descr, false, shape, &[0u8; 64]);
1593            let len = bytes.len();
1594            open_in(Arc::new(Bytes::Owned(bytes)), 0, len, "x")
1595                .err()
1596                .unwrap()
1597        };
1598        assert!(refused("'|O'", "(1,)").contains("does not unpickle"));
1599        assert!(refused("'<f8'", "(1, 1, 1)").contains("shape (1, 1, 1)"));
1600        assert!(refused("'<f16'", "(1,)").contains("not read"));
1601    }
1602
1603    #[test]
1604    fn a_short_file_keeps_the_rows_it_has() {
1605        let data: Vec<u8> = [1i16, 2, 3].iter().flat_map(|v| v.to_le_bytes()).collect();
1606        let bytes = npy("'<i2'", false, "(10,)", &data);
1607        let len = bytes.len();
1608        let array = open_in(Arc::new(Bytes::Owned(bytes)), 0, len, "x").unwrap();
1609        assert_eq!(array.records.rows(), 3);
1610        assert_eq!(array.notes.len(), 1);
1611    }
1612}