mod elements;
pub mod header;
use self::header::{
FormatHeaderError, Header, Layout, ParseHeaderError, ReadHeaderError, WriteHeaderError,
};
use ndarray::prelude::*;
use ndarray::{Data, DataOwned, IntoDimension};
use py_literal::Value as PyValue;
use std::convert::TryInto;
use std::error::Error;
use std::fmt;
use std::fs::File;
use std::io::{self, BufWriter, Seek};
use std::mem;
pub fn read_npy<P, T>(path: P) -> Result<T, ReadNpyError>
where
P: AsRef<std::path::Path>,
T: ReadNpyExt,
{
T::read_npy(File::open(path)?)
}
pub fn write_npy<P, T>(path: P, array: &T) -> Result<(), WriteNpyError>
where
P: AsRef<std::path::Path>,
T: WriteNpyExt + ?Sized,
{
array.write_npy(BufWriter::new(File::create(path)?))
}
pub fn create_new_npy<P, T>(path: P, array: &T) -> Result<(), WriteNpyError>
where
P: AsRef<std::path::Path>,
T: WriteNpyExt + ?Sized,
{
array.write_npy(BufWriter::new(File::create_new(path)?))
}
pub fn write_zeroed_npy<A, D>(mut file: &File, shape: D) -> Result<(), WriteNpyError>
where
A: WritableElement,
D: IntoDimension,
{
let dim = shape.into_dimension();
let data_bytes_len: u64 = dim
.size_checked()
.expect("overflow computing number of elements")
.checked_mul(mem::size_of::<A>())
.expect("overflow computing length of data")
.try_into()
.expect("overflow converting length of data to u64");
Header {
type_descriptor: A::type_descriptor(),
layout: Layout::Standard,
shape: dim.as_array_view().to_vec(),
}
.write(file)?;
let current_offset = file.stream_position()?;
file.set_len(current_offset)?;
file.set_len(
current_offset
.checked_add(data_bytes_len)
.expect("overflow computing file length"),
)?;
Ok(())
}
#[derive(Debug)]
pub enum WriteDataError {
Io(io::Error),
FormatData(Box<dyn Error + Send + Sync + 'static>),
}
impl Error for WriteDataError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
match self {
WriteDataError::Io(err) => Some(err),
WriteDataError::FormatData(err) => Some(&**err),
}
}
}
impl fmt::Display for WriteDataError {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
WriteDataError::Io(err) => write!(f, "I/O error: {}", err),
WriteDataError::FormatData(err) => write!(f, "error formatting data: {}", err),
}
}
}
impl From<io::Error> for WriteDataError {
fn from(err: io::Error) -> WriteDataError {
WriteDataError::Io(err)
}
}
pub trait WritableElement: Sized {
fn type_descriptor() -> PyValue;
fn write<W: io::Write>(&self, writer: W) -> Result<(), WriteDataError>;
fn write_slice<W: io::Write>(slice: &[Self], writer: W) -> Result<(), WriteDataError>;
}
#[derive(Debug)]
pub enum WriteNpyError {
Io(io::Error),
FormatHeader(FormatHeaderError),
FormatData(Box<dyn Error + Send + Sync + 'static>),
}
impl Error for WriteNpyError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
match self {
WriteNpyError::Io(err) => Some(err),
WriteNpyError::FormatHeader(err) => Some(err),
WriteNpyError::FormatData(err) => Some(&**err),
}
}
}
impl fmt::Display for WriteNpyError {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
WriteNpyError::Io(err) => write!(f, "I/O error: {}", err),
WriteNpyError::FormatHeader(err) => write!(f, "error formatting header: {}", err),
WriteNpyError::FormatData(err) => write!(f, "error formatting data: {}", err),
}
}
}
impl From<io::Error> for WriteNpyError {
fn from(err: io::Error) -> WriteNpyError {
WriteNpyError::Io(err)
}
}
impl From<WriteHeaderError> for WriteNpyError {
fn from(err: WriteHeaderError) -> WriteNpyError {
match err {
WriteHeaderError::Io(err) => WriteNpyError::Io(err),
WriteHeaderError::Format(err) => WriteNpyError::FormatHeader(err),
}
}
}
impl From<FormatHeaderError> for WriteNpyError {
fn from(err: FormatHeaderError) -> WriteNpyError {
WriteNpyError::FormatHeader(err)
}
}
impl From<WriteDataError> for WriteNpyError {
fn from(err: WriteDataError) -> WriteNpyError {
match err {
WriteDataError::Io(err) => WriteNpyError::Io(err),
WriteDataError::FormatData(err) => WriteNpyError::FormatData(err),
}
}
}
pub trait WriteNpyExt {
fn write_npy<W: io::Write>(&self, writer: W) -> Result<(), WriteNpyError>;
}
impl<A, D> WriteNpyExt for ArrayRef<A, D>
where
A: WritableElement,
D: Dimension,
{
fn write_npy<W: io::Write>(&self, mut writer: W) -> Result<(), WriteNpyError> {
let write_contiguous = |mut writer: W, layout: Layout| {
Header {
type_descriptor: A::type_descriptor(),
layout,
shape: self.shape().to_owned(),
}
.write(&mut writer)?;
A::write_slice(self.as_slice_memory_order().unwrap(), &mut writer)?;
writer.flush()?;
Ok(())
};
if self.is_standard_layout() {
write_contiguous(writer, Layout::Standard)
} else if self.view().reversed_axes().is_standard_layout() {
write_contiguous(writer, Layout::Fortran)
} else {
Header {
type_descriptor: A::type_descriptor(),
layout: Layout::Standard,
shape: self.shape().to_owned(),
}
.write(&mut writer)?;
for elem in self.iter() {
elem.write(&mut writer)?;
}
writer.flush()?;
Ok(())
}
}
}
impl<A, S, D> WriteNpyExt for ArrayBase<S, D>
where
A: WritableElement,
S: Data<Elem = A>,
D: Dimension,
{
fn write_npy<W: io::Write>(&self, writer: W) -> Result<(), WriteNpyError> {
let arr: &ArrayRef<A, D> = self;
arr.write_npy(writer)
}
}
#[derive(Debug)]
pub enum ReadDataError {
Io(io::Error),
WrongDescriptor(PyValue),
MissingData,
ExtraBytes(usize),
ParseData(Box<dyn Error + Send + Sync + 'static>),
}
impl Error for ReadDataError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
match self {
ReadDataError::Io(err) => Some(err),
ReadDataError::WrongDescriptor(_) => None,
ReadDataError::MissingData => None,
ReadDataError::ExtraBytes(_) => None,
ReadDataError::ParseData(err) => Some(&**err),
}
}
}
impl fmt::Display for ReadDataError {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
ReadDataError::Io(err) => write!(f, "I/O error: {}", err),
ReadDataError::WrongDescriptor(desc) => {
write!(f, "incorrect descriptor ({}) for this type", desc)
}
ReadDataError::MissingData => write!(f, "reached EOF before reading all data"),
ReadDataError::ExtraBytes(num_extra_bytes) => {
write!(f, "file had {} extra bytes before EOF", num_extra_bytes)
}
ReadDataError::ParseData(err) => write!(f, "error parsing data: {}", err),
}
}
}
impl From<io::Error> for ReadDataError {
fn from(err: io::Error) -> ReadDataError {
if err.kind() == io::ErrorKind::UnexpectedEof {
ReadDataError::MissingData
} else {
ReadDataError::Io(err)
}
}
}
pub trait ReadableElement: Sized {
fn read_to_end_exact_vec<R: io::Read>(
reader: R,
type_desc: &PyValue,
len: usize,
) -> Result<Vec<Self>, ReadDataError>;
}
#[derive(Debug)]
pub enum ReadNpyError {
Io(io::Error),
ParseHeader(ParseHeaderError),
ParseData(Box<dyn Error + Send + Sync + 'static>),
LengthOverflow,
WrongNdim(Option<usize>, usize),
WrongDescriptor(PyValue),
MissingData,
ExtraBytes(usize),
}
impl Error for ReadNpyError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
match self {
ReadNpyError::Io(err) => Some(err),
ReadNpyError::ParseHeader(err) => Some(err),
ReadNpyError::ParseData(err) => Some(&**err),
ReadNpyError::LengthOverflow => None,
ReadNpyError::WrongNdim(_, _) => None,
ReadNpyError::WrongDescriptor(_) => None,
ReadNpyError::MissingData => None,
ReadNpyError::ExtraBytes(_) => None,
}
}
}
impl fmt::Display for ReadNpyError {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
ReadNpyError::Io(err) => write!(f, "I/O error: {}", err),
ReadNpyError::ParseHeader(err) => write!(f, "error parsing header: {}", err),
ReadNpyError::ParseData(err) => write!(f, "error parsing data: {}", err),
ReadNpyError::LengthOverflow => write!(f, "overflow computing length from shape"),
ReadNpyError::WrongNdim(expected, actual) => write!(
f,
"ndim {} of array did not match Dimension type with NDIM = {:?}",
actual, expected
),
ReadNpyError::WrongDescriptor(desc) => {
write!(f, "incorrect descriptor ({}) for this type", desc)
}
ReadNpyError::MissingData => write!(f, "reached EOF before reading all data"),
ReadNpyError::ExtraBytes(num_extra_bytes) => {
write!(f, "file had {} extra bytes before EOF", num_extra_bytes)
}
}
}
}
impl From<io::Error> for ReadNpyError {
fn from(err: io::Error) -> ReadNpyError {
ReadNpyError::Io(err)
}
}
impl From<ReadHeaderError> for ReadNpyError {
fn from(err: ReadHeaderError) -> ReadNpyError {
match err {
ReadHeaderError::Io(err) => ReadNpyError::Io(err),
ReadHeaderError::Parse(err) => ReadNpyError::ParseHeader(err),
}
}
}
impl From<ParseHeaderError> for ReadNpyError {
fn from(err: ParseHeaderError) -> ReadNpyError {
ReadNpyError::ParseHeader(err)
}
}
impl From<ReadDataError> for ReadNpyError {
fn from(err: ReadDataError) -> ReadNpyError {
match err {
ReadDataError::Io(err) => ReadNpyError::Io(err),
ReadDataError::WrongDescriptor(desc) => ReadNpyError::WrongDescriptor(desc),
ReadDataError::MissingData => ReadNpyError::MissingData,
ReadDataError::ExtraBytes(nbytes) => ReadNpyError::ExtraBytes(nbytes),
ReadDataError::ParseData(err) => ReadNpyError::ParseData(err),
}
}
}
pub trait ReadNpyExt: Sized {
fn read_npy<R: io::Read>(reader: R) -> Result<Self, ReadNpyError>;
}
impl<A, S, D> ReadNpyExt for ArrayBase<S, D>
where
A: ReadableElement,
S: DataOwned<Elem = A>,
D: Dimension,
{
fn read_npy<R: io::Read>(mut reader: R) -> Result<Self, ReadNpyError> {
let header = Header::from_reader(&mut reader)?;
let shape = header.shape.into_dimension();
let ndim = shape.ndim();
let len = shape_length_checked::<A>(&shape).ok_or(ReadNpyError::LengthOverflow)?;
let data = A::read_to_end_exact_vec(&mut reader, &header.type_descriptor, len)?;
ArrayBase::from_shape_vec(shape.set_f(header.layout.is_fortran()), data)
.unwrap()
.into_dimensionality()
.map_err(|_| ReadNpyError::WrongNdim(D::NDIM, ndim))
}
}
#[derive(Debug)]
#[non_exhaustive]
pub enum ViewNpyError {
Io(io::Error),
ParseHeader(ParseHeaderError),
InvalidData(Box<dyn Error + Send + Sync + 'static>),
LengthOverflow,
WrongNdim(Option<usize>, usize),
WrongDescriptor(PyValue),
NonNativeEndian,
MisalignedData,
MissingBytes(usize),
ExtraBytes(usize),
}
impl Error for ViewNpyError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
match self {
ViewNpyError::Io(err) => Some(err),
ViewNpyError::ParseHeader(err) => Some(err),
ViewNpyError::InvalidData(err) => Some(&**err),
ViewNpyError::LengthOverflow => None,
ViewNpyError::WrongNdim(_, _) => None,
ViewNpyError::WrongDescriptor(_) => None,
ViewNpyError::NonNativeEndian => None,
ViewNpyError::MisalignedData => None,
ViewNpyError::MissingBytes(_) => None,
ViewNpyError::ExtraBytes(_) => None,
}
}
}
impl fmt::Display for ViewNpyError {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
ViewNpyError::Io(err) => write!(f, "I/O error: {}", err),
ViewNpyError::ParseHeader(err) => write!(f, "error parsing header: {}", err),
ViewNpyError::InvalidData(err) => write!(f, "invalid data for element type: {}", err),
ViewNpyError::LengthOverflow => write!(f, "overflow computing length from shape"),
ViewNpyError::WrongNdim(expected, actual) => write!(
f,
"ndim {} of array did not match Dimension type with NDIM = {:?}",
actual, expected
),
ViewNpyError::WrongDescriptor(desc) => {
write!(f, "incorrect descriptor ({}) for this type", desc)
}
ViewNpyError::NonNativeEndian => {
write!(f, "descriptor does not match native endianness")
}
ViewNpyError::MisalignedData => write!(
f,
"start of data is not properly aligned for the element type"
),
ViewNpyError::MissingBytes(num_missing_bytes) => write!(
f,
"missing {} bytes of data specified in header",
num_missing_bytes
),
ViewNpyError::ExtraBytes(num_extra_bytes) => {
write!(f, "file had {} extra bytes before EOF", num_extra_bytes)
}
}
}
}
impl From<ReadHeaderError> for ViewNpyError {
fn from(err: ReadHeaderError) -> ViewNpyError {
match err {
ReadHeaderError::Io(err) => ViewNpyError::Io(err),
ReadHeaderError::Parse(err) => ViewNpyError::ParseHeader(err),
}
}
}
impl From<ParseHeaderError> for ViewNpyError {
fn from(err: ParseHeaderError) -> ViewNpyError {
ViewNpyError::ParseHeader(err)
}
}
impl From<ViewDataError> for ViewNpyError {
fn from(err: ViewDataError) -> ViewNpyError {
match err {
ViewDataError::WrongDescriptor(desc) => ViewNpyError::WrongDescriptor(desc),
ViewDataError::NonNativeEndian => ViewNpyError::NonNativeEndian,
ViewDataError::Misaligned => ViewNpyError::MisalignedData,
ViewDataError::MissingBytes(nbytes) => ViewNpyError::MissingBytes(nbytes),
ViewDataError::ExtraBytes(nbytes) => ViewNpyError::ExtraBytes(nbytes),
ViewDataError::InvalidData(err) => ViewNpyError::InvalidData(err),
}
}
}
pub trait ViewNpyExt<'a>: Sized {
fn view_npy(buf: &'a [u8]) -> Result<Self, ViewNpyError>;
}
pub trait ViewMutNpyExt<'a>: Sized {
fn view_mut_npy(buf: &'a mut [u8]) -> Result<Self, ViewNpyError>;
}
impl<'a, A, D> ViewNpyExt<'a> for ArrayView<'a, A, D>
where
A: ViewElement,
D: Dimension,
{
fn view_npy(buf: &'a [u8]) -> Result<Self, ViewNpyError> {
let mut reader = buf;
let header = Header::from_reader(&mut reader)?;
let shape = header.shape.into_dimension();
let ndim = shape.ndim();
let len = shape_length_checked::<A>(&shape).ok_or(ViewNpyError::LengthOverflow)?;
let data = A::bytes_as_slice(reader, &header.type_descriptor, len)?;
ArrayView::from_shape(shape.set_f(header.layout.is_fortran()), data)
.unwrap()
.into_dimensionality()
.map_err(|_| ViewNpyError::WrongNdim(D::NDIM, ndim))
}
}
impl<'a, A, D> ViewMutNpyExt<'a> for ArrayViewMut<'a, A, D>
where
A: ViewMutElement,
D: Dimension,
{
fn view_mut_npy(buf: &'a mut [u8]) -> Result<Self, ViewNpyError> {
let mut reader = &*buf;
let header = Header::from_reader(&mut reader)?;
let shape = header.shape.into_dimension();
let ndim = shape.ndim();
let len = shape_length_checked::<A>(&shape).ok_or(ViewNpyError::LengthOverflow)?;
let mid = buf.len() - reader.len();
let data = A::bytes_as_mut_slice(&mut buf[mid..], &header.type_descriptor, len)?;
ArrayViewMut::from_shape(shape.set_f(header.layout.is_fortran()), data)
.unwrap()
.into_dimensionality()
.map_err(|_| ViewNpyError::WrongNdim(D::NDIM, ndim))
}
}
#[derive(Debug)]
#[non_exhaustive]
pub enum ViewDataError {
WrongDescriptor(PyValue),
NonNativeEndian,
Misaligned,
MissingBytes(usize),
ExtraBytes(usize),
InvalidData(Box<dyn Error + Send + Sync + 'static>),
}
impl Error for ViewDataError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
match self {
ViewDataError::WrongDescriptor(_) => None,
ViewDataError::NonNativeEndian => None,
ViewDataError::Misaligned => None,
ViewDataError::MissingBytes(_) => None,
ViewDataError::ExtraBytes(_) => None,
ViewDataError::InvalidData(err) => Some(&**err),
}
}
}
impl fmt::Display for ViewDataError {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
ViewDataError::WrongDescriptor(desc) => {
write!(f, "incorrect descriptor ({}) for this type", desc)
}
ViewDataError::NonNativeEndian => {
write!(f, "descriptor does not match native endianness")
}
ViewDataError::Misaligned => write!(
f,
"start of data is not properly aligned for the element type"
),
ViewDataError::MissingBytes(num_missing_bytes) => write!(
f,
"missing {} bytes of data specified in header",
num_missing_bytes
),
ViewDataError::ExtraBytes(num_extra_bytes) => {
write!(f, "file had {} extra bytes before EOF", num_extra_bytes)
}
ViewDataError::InvalidData(err) => write!(f, "invalid data for element type: {}", err),
}
}
}
pub trait ViewElement: Sized {
fn bytes_as_slice<'a>(
bytes: &'a [u8],
type_desc: &PyValue,
len: usize,
) -> Result<&'a [Self], ViewDataError>;
}
pub trait ViewMutElement: Sized {
fn bytes_as_mut_slice<'a>(
bytes: &'a mut [u8],
type_desc: &PyValue,
len: usize,
) -> Result<&'a mut [Self], ViewDataError>;
}
fn shape_length_checked<T>(shape: &IxDyn) -> Option<usize> {
let len = shape.size_checked()?;
if len > isize::MAX as usize {
return None;
}
let bytes = len.checked_mul(mem::size_of::<T>())?;
if bytes > isize::MAX as usize {
return None;
}
Some(len)
}