Skip to main content

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