diff --git a/Cargo.toml b/Cargo.toml index 6eee20a..e6e126c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,8 +1,8 @@ [package] name = "darn-dmap" -version = "0.8.0" +version = "0.8.1" edition = "2021" -rust-version = "1.63.0" +rust-version = "1.85.1" authors = ["Remington Rohel"] description = "SuperDARN DMAP file format I/O" repository = "https://github.com/SuperDARNCanada/dmap" @@ -16,12 +16,17 @@ include = ["src/**/*.rs"] [lib] name = "dmap" -# "cdylib" is necessary to produce a shared library for Python to import from. +# cdylib is needed when building the optional Python extension module. crate-type = ["cdylib", "rlib"] +[features] +default = [] +python = ["dep:pyo3", "dep:numpy"] + [dependencies] -pyo3 = { version = "0.26.0", features = ["extension-module", "indexmap", "abi3-py38"] } -numpy = "0.26.0" +ndarray = "0.16" +pyo3 = { version = "0.26.0", optional = true, features = ["extension-module", "indexmap", "abi3-py38"] } +numpy = { version = "0.26.0", optional = true } indexmap = "2.3.0" itertools = "0.13.0" rayon = "1.10.0" diff --git a/pyproject.toml b/pyproject.toml index c90f855..53be8dc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "darn-dmap" -version = "0.8.0" +version = "0.8.1" requires-python = ">=3.8" authors = [ { name = "Remington Rohel" } @@ -26,6 +26,7 @@ python-source = "python" module-name = "dmap.dmap_rs" bindings = "pyo3" profile = "release" +features = ["python"] compatibility = "manylinux2014" auditwheel = "repair" strip = true diff --git a/src/convenience.rs b/src/convenience.rs index d4221c1..c6263c2 100644 --- a/src/convenience.rs +++ b/src/convenience.rs @@ -18,5 +18,3 @@ pub(crate) fn split_results(dmap_results: Vec>) -> (Vec, V (ok_inners, error_inners, bad_indices) } - - diff --git a/src/error.rs b/src/error.rs index c2b6493..844187f 100644 --- a/src/error.rs +++ b/src/error.rs @@ -1,5 +1,7 @@ //! Error type for `dmap`. +#[cfg(feature = "python")] use pyo3::exceptions::{PyIOError, PyValueError}; +#[cfg(feature = "python")] use pyo3::PyErr; use thiserror::Error; @@ -38,7 +40,7 @@ pub enum DmapError { /// Bytes cannot be interpreted as a DMAP field. #[error("{0}")] InvalidField(String), - + /// Index out of bounds. #[error("{0}")] InvalidIndex(i32), @@ -48,6 +50,7 @@ pub enum DmapError { BadRecords(Vec, String), } +#[cfg(feature = "python")] impl From for PyErr { fn from(value: DmapError) -> Self { let msg = value.to_string(); diff --git a/src/io.rs b/src/io.rs index 94106f5..6ad8bc7 100644 --- a/src/io.rs +++ b/src/io.rs @@ -1,6 +1,7 @@ //! Utility functions for file operations. use crate::compression::{compress_bz2, detect_bz2}; +use crate::parser::Parser; use crate::types::DmapType; use crate::DmapError; use bzip2::read::BzDecoder; @@ -8,7 +9,6 @@ use std::ffi::OsStr; use std::fs::{File, OpenOptions}; use std::io::{Read, Write}; use std::path::Path; -use crate::parser::Parser; /// Write bytes to file. /// @@ -47,7 +47,9 @@ pub(crate) fn write_bytes_bz2( } /// Set up the stream for reading. Autodetects and decompresses BZ2. -pub(crate) fn create_stream<'b>(dmap_data: &'b mut impl Read) -> Result, DmapError> { +pub(crate) fn create_stream<'b>( + dmap_data: &'b mut impl Read, +) -> Result, DmapError> { let (is_bz2, chunk) = detect_bz2(dmap_data)?; if is_bz2 { Ok(Box::new(BzDecoder::new(chunk))) @@ -57,9 +59,7 @@ pub(crate) fn create_stream<'b>(dmap_data: &'b mut impl Read) -> Result Result, DmapError> { +pub(crate) fn split_into_slices(mut dmap_data: impl Read) -> Result, DmapError> { let mut buffer: Vec = vec![]; create_stream(&mut dmap_data)?.read_to_end(&mut buffer)?; @@ -107,8 +107,8 @@ pub(crate) fn slice_stream_lax(buffer: Vec) -> (Vec, Vec, Opt let mut rec_starts = vec![]; while ((rec_start + 2 * i32::size()) as u64) < buffer.len() as u64 { - rec_size = i32::from_le_bytes(buffer[rec_start + 4..rec_start + 8].try_into().unwrap()) - as usize; // advance 4 bytes, skipping the "code" field + rec_size = + i32::from_le_bytes(buffer[rec_start + 4..rec_start + 8].try_into().unwrap()) as usize; // advance 4 bytes, skipping the "code" field rec_end = rec_start + rec_size; // error-checking the size is conducted in Self::parse_record() if rec_end > buffer.len() || rec_size == 0 { bad_byte = Some(rec_start); diff --git a/src/lib.rs b/src/lib.rs index fb8c42b..823043c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -92,9 +92,15 @@ pub(crate) mod convenience; pub mod error; pub mod formats; pub(crate) mod io; +mod parser; pub mod record; pub mod types; -mod parser; + +#[cfg(feature = "python")] +mod python_types; + +#[cfg(feature = "python")] +mod python_bindings; pub use crate::error::DmapError; pub use crate::formats::dmap::DmapRecord; @@ -105,478 +111,3 @@ pub use crate::formats::map::MapRecord; pub use crate::formats::rawacf::RawacfRecord; pub use crate::formats::snd::SndRecord; pub use crate::record::Record; -use crate::types::DmapField; -use indexmap::IndexMap; -use paste::paste; -use pyo3::prelude::*; -use pyo3::types::PyBytes; -use std::path::{Path, PathBuf}; - -/// This macro generates a function for attempting to convert `Vec` to `Vec<$type>` and write it to file. -macro_rules! write_rust { - ($type:ident) => { - paste! { - #[doc = "Attempts to convert `recs` to `" $type:camel Record "` then append to `outfile`."] - #[doc = ""] - #[doc = "# Errors"] - #[doc = "if any of the `IndexMap`s are unable to be interpreted as a `" $type:camel Record "`, or there is an issue writing to file."] - pub fn [< try_write_ $type >]>( - recs: Vec>, - outfile: P, - bz2: bool, - ) -> Result<(), DmapError> { - let bytes = [< $type:camel Record >]::try_into_bytes(recs)?; - crate::io::bytes_to_file(bytes, outfile, bz2).map_err(DmapError::from) - } - } - } -} - -write_rust!(iqdat); -write_rust!(rawacf); -write_rust!(fitacf); -write_rust!(grid); -write_rust!(map); -write_rust!(snd); -write_rust!(dmap); - -/// Creates functions for reading DMAP files for the Python API. -/// -/// Generates six functions: -/// * `read_[name]` - reads a file, raising an error on a corrupted file. -/// * `read_[name]_lax` - reads a file, returning the records and the byte where corruption starts, if corrupted. -/// * `read_[name]_bytes` - reads from bytes, similar to `read_[name]`. -/// * `read_[name]_bytes_lax` - reads from bytes, similar to `read_[name]_lax`. -/// * `read_[name]_by_index` - reads only the nth record(s) from file. -/// * `read_[name]_by_index_lax` - reads only the nth record(s) from file in the same manner as `read_[name]_lax`. -/// * `read_[name]_metadata` - reads only the metadata from records in a file. -/// reading, respectively. -macro_rules! read_py { - ( - $name:ident, - $py_name:literal, - $lax_name:literal, - $bytes_name:literal, - $lax_bytes_name:literal, - $sniff_name:literal, - $sniff_name_lax:literal, - $sniff_bytes:literal, - $sniff_bytes_lax:literal, - $metadata_name:literal, - $sniff_metadata:literal - ) => { - paste! { - #[doc = "Reads a `" $name:upper "` file, returning a list of dictionaries containing the fields." ] - #[pyfunction] - #[pyo3(name = $py_name)] - #[pyo3(text_signature = "(infile: str, /)")] - fn [< read_ $name _py >](infile: PathBuf) -> PyResult>> { - Ok([< $name:camel Record >]::read_file(&infile) - .map_err(PyErr::from)? - .into_iter() - .map(|rec| rec.inner()) - .collect() - ) - } - - #[doc = "Reads a `" $name:upper "` file, returning a tuple of" ] - #[doc = "(list of dictionaries containing the fields, byte where first corrupted record starts). "] - #[pyfunction] - #[pyo3(name = $lax_name)] - #[pyo3(text_signature = "(infile: str, /)")] - fn [< read_ $name _lax_py >]( - infile: PathBuf, - ) -> PyResult<(Vec>, Option)> { - let result = [< $name:camel Record >]::read_file_lax(&infile).map_err(PyErr::from)?; - Ok(( - result.0.into_iter().map(|rec| rec.inner()).collect(), - result.1, - )) - } - - #[doc = "Read in `" $name:upper "` records from bytes, returning `List[Dict]` of the records." ] - #[pyfunction] - #[pyo3(name = $bytes_name)] - #[pyo3(text_signature = "(buf: bytes, /)")] - fn [< read_ $name _bytes_py >](bytes: &[u8]) -> PyResult>> { - Ok([< $name:camel Record >]::read_records(bytes)? - .into_iter() - .map(|rec| rec.inner()) - .collect() - ) - } - - #[doc = "Reads a `" $name:upper "` file, returning a tuple of" ] - #[doc = "(list of dictionaries containing the fields, byte where first corrupted record starts). "] - #[pyfunction] - #[pyo3(name = $lax_bytes_name)] - #[pyo3(text_signature = "(buf: bytes, /)")] - fn [< read_ $name _bytes_lax_py >]( - bytes: &[u8], - ) -> PyResult<(Vec>, Option)> { - let result = [< $name:camel Record >]::read_records_lax(bytes).map_err(PyErr::from)?; - Ok(( - result.0.into_iter().map(|rec| rec.inner()).collect(), - result.1, - )) - } - - #[doc = "Reads a `" $name:upper "` file, returning the `nth` record(s)." ] - #[pyfunction] - #[pyo3(name = $sniff_name)] - #[pyo3(text_signature = "(infile: str, indices: tuple[int], /)")] - fn [< read_ $name _by_indices_py >](infile: PathBuf, indices: Vec) -> PyResult>> { - Ok([< $name:camel Record >]::read_file_by_indices(&infile, &indices) - .map_err(PyErr::from)? - .into_iter() - .map(|rec| rec.inner()) - .collect() - ) - } - - #[doc = "Reads a `" $name:upper "` file, returning the `nth` record(s), and the byte index where corruption starts, if applicable" ] - #[pyfunction] - #[pyo3(name = $sniff_name_lax)] - #[pyo3(text_signature = "(infile: str, indices: tuple[int], /)")] - fn [< read_ $name _by_indices_lax_py >]( - infile: PathBuf, indices: Vec - ) -> PyResult<(Vec>, Option)> { - let result = [< $name:camel Record >]::read_file_by_indices_lax(infile, &indices).map_err(PyErr::from)?; - Ok(( - result.0.into_iter().map(|rec| rec.inner()).collect(), - result.1, - )) - } - - #[doc = "Reads a `" $name:upper "` buffer, returning the `nth` record(s)." ] - #[pyfunction] - #[pyo3(name = $sniff_bytes)] - #[pyo3(text_signature = "(buf: bytes, indices: tuple[int], /)")] - fn [< read_ $name _bytes_by_indices_py >](buf: &[u8], indices: Vec) -> PyResult>> { - Ok([< $name:camel Record >]::read_nth_records(buf, &indices) - .map_err(PyErr::from)? - .into_iter() - .map(|rec| rec.inner()) - .collect() - ) - } - - #[doc = "Reads a `" $name:upper "` buffer, returning the `nth` record(s) and the byte index where record corruption starts, if applicable." ] - #[pyfunction] - #[pyo3(name = $sniff_bytes_lax)] - #[pyo3(text_signature = "(buf: bytes, indices: tuple[int], /)")] - fn [< read_ $name _bytes_by_indices_lax_py >]( - buf: &[u8], indices: Vec - ) -> PyResult<(Vec>, Option)> { - let result = [< $name:camel Record >]::read_nth_records_lax(buf, &indices).map_err(PyErr::from)?; - Ok(( - result.0.into_iter().map(|rec| rec.inner()).collect(), - result.1, - )) - } - - #[doc = "Reads a `" $name:upper "` file, returning a list of dictionaries containing the only the metadata fields." ] - #[pyfunction] - #[pyo3(name = $metadata_name)] - #[pyo3(text_signature = "(infile: str, /)")] - fn [< read_ $name _metadata_py >](infile: PathBuf) -> PyResult>> { - Ok([< $name:camel Record >]::read_file_metadata(&infile) - .map_err(PyErr::from)? - ) - } - - #[doc = "Reads a `" $name:upper "` file, returning the `nth` records' metadata fields." ] - #[pyfunction] - #[pyo3(name = $sniff_metadata)] - #[pyo3(text_signature = "(infile: str, indices: tuple[int], /)")] - fn [< read_ $name _metadata_by_indices_py >](infile: PathBuf, indices: Vec) -> PyResult>> { - Ok([< $name:camel Record >]::read_file_metadata_by_indices(&infile, &indices) - .map_err(PyErr::from)? - ) - } - } - } -} - -read_py!( - iqdat, - "read_iqdat", - "read_iqdat_lax", - "read_iqdat_bytes", - "read_iqdat_bytes_lax", - "read_iqdat_by_indices", - "read_iqdat_by_indices_lax", - "read_iqdat_by_indices_bytes", - "read_iqdat_by_indices_bytes_lax", - "read_iqdat_metadata", - "read_iqdat_metadata_by_indices" -); -read_py!( - rawacf, - "read_rawacf", - "read_rawacf_lax", - "read_rawacf_bytes", - "read_rawacf_bytes_lax", - "read_rawacf_by_indices", - "read_rawacf_by_indices_lax", - "read_rawacf_by_indices_bytes", - "read_rawacf_by_indices_bytes_lax", - "read_rawacf_metadata", - "read_rawacf_metadata_by_indices" -); -read_py!( - fitacf, - "read_fitacf", - "read_fitacf_lax", - "read_fitacf_bytes", - "read_fitacf_bytes_lax", - "read_fitacf_by_indices", - "read_fitacf_by_indices_lax", - "read_fitacf_by_indices_bytes", - "read_fitacf_by_indices_bytes_lax", - "read_fitacf_metadata", - "read_fitacf_metadata_by_indices" -); -read_py!( - grid, - "read_grid", - "read_grid_lax", - "read_grid_bytes", - "read_grid_bytes_lax", - "read_grid_by_indices", - "read_grid_by_indices_lax", - "read_grid_by_indices_bytes", - "read_grid_by_indices_bytes_lax", - "read_grid_metadata", - "read_grid_metadata_by_indices" -); -read_py!( - map, - "read_map", - "read_map_lax", - "read_map_bytes", - "read_map_bytes_lax", - "read_map_by_indices", - "read_map_by_indices_lax", - "read_map_by_indices_bytes", - "read_map_by_indices_bytes_lax", - "read_map_metadata", - "read_map_metadata_by_indices" -); -read_py!( - snd, - "read_snd", - "read_snd_lax", - "read_snd_bytes", - "read_snd_bytes_lax", - "read_snd_by_indices", - "read_snd_by_indices_lax", - "read_snd_by_indices_bytes", - "read_snd_by_indices_bytes_lax", - "read_snd_metadata", - "read_snd_metadata_by_indices" -); -read_py!( - dmap, - "read_dmap", - "read_dmap_lax", - "read_dmap_bytes", - "read_dmap_bytes_lax", - "read_dmap_by_indices", - "read_dmap_by_indices_lax", - "read_dmap_by_indices_bytes", - "read_dmap_by_indices_bytes_lax", - "read_dmap_metadata", - "read_dmap_metadata_by_indices" -); - -/// Checks that a list of dictionaries contains DMAP records, then appends to outfile. -/// -/// **NOTE:** No type checking is done, so the fields may not be written as the expected -/// DMAP type, e.g. `stid` might be written one byte instead of two as this function -/// does not know that typically `stid` is two bytes. -#[pyfunction] -#[pyo3(name = "write_dmap")] -#[pyo3(signature = (recs, outfile, /, bz2))] -#[pyo3(text_signature = "(recs: list[dict], outfile: str, /, bz2: bool = False)")] -fn write_dmap_py( - recs: Vec>, - outfile: PathBuf, - bz2: bool, -) -> PyResult<()> { - try_write_dmap(recs, &outfile, bz2).map_err(PyErr::from) -} - -/// Checks that a list of dictionaries contains valid DMAP records, then converts them to bytes. -/// Returns `list[bytes]`, one entry per record. -/// -/// **NOTE:** No type checking is done, so the fields may not be written as the expected -/// DMAP type, e.g. `stid` might be written one byte instead of two as this function -/// does not know that typically `stid` is two bytes. -#[pyfunction] -#[pyo3(name = "write_dmap_bytes")] -#[pyo3(signature = (recs, /, bz2))] -#[pyo3(text_signature = "(recs: list[dict], /, bz2: bool = False)")] -fn write_dmap_bytes_py( - py: Python, - recs: Vec>, - bz2: bool, -) -> PyResult> { - let mut bytes = DmapRecord::try_into_bytes(recs).map_err(PyErr::from)?; - if bz2 { - bytes = compression::compress_bz2(&bytes).map_err(PyErr::from)?; - } - Ok(PyBytes::new(py, &bytes).into()) -} - -/// Generates functions exposed to the Python API for writing specific file types. -macro_rules! write_py { - ($name:ident, $fn_name:literal, $bytes_name:literal) => { - paste! { - #[doc = "Checks that a list of dictionaries contains valid `" $name:upper "` records, then appends to outfile." ] - #[pyfunction] - #[pyo3(name = $fn_name)] - #[pyo3(signature = (recs, outfile, /, bz2))] - #[pyo3(text_signature = "(recs: list[dict], outfile: str, /, bz2: bool = False)")] - fn [< write_ $name _py >](recs: Vec>, outfile: PathBuf, bz2: bool) -> PyResult<()> { - [< try_write_ $name >](recs, &outfile, bz2).map_err(PyErr::from) - } - - #[doc = "Checks that a list of dictionaries contains valid `" $name:upper "` records, then converts them to bytes." ] - #[doc = "Returns `list[bytes]`, one entry per record." ] - #[pyfunction] - #[pyo3(name = $bytes_name)] - #[pyo3(signature = (recs, /, bz2))] - #[pyo3(text_signature = "(recs: list[dict], /, bz2: bool = False)")] - fn [< write_ $name _bytes_py >](py: Python, recs: Vec>, bz2: bool) -> PyResult> { - let mut bytes = [< $name:camel Record >]::try_into_bytes(recs).map_err(PyErr::from)?; - if bz2 { - bytes = compression::compress_bz2(&bytes).map_err(PyErr::from)?; - } - Ok(PyBytes::new(py, &bytes).into()) - } - } - } -} - -// **NOTE** dmap type not included in this list, since it has a more descriptive docstring. -write_py!(iqdat, "write_iqdat", "write_iqdat_bytes"); -write_py!(rawacf, "write_rawacf", "write_rawacf_bytes"); -write_py!(fitacf, "write_fitacf", "write_fitacf_bytes"); -write_py!(grid, "write_grid", "write_grid_bytes"); -write_py!(map, "write_map", "write_map_bytes"); -write_py!(snd, "write_snd", "write_snd_bytes"); - -/// Functions for SuperDARN DMAP file format I/O. -#[pymodule] -fn dmap_rs(m: &Bound<'_, PyModule>) -> PyResult<()> { - // Strict read functions - m.add_function(wrap_pyfunction!(read_dmap_py, m)?)?; - m.add_function(wrap_pyfunction!(read_iqdat_py, m)?)?; - m.add_function(wrap_pyfunction!(read_rawacf_py, m)?)?; - m.add_function(wrap_pyfunction!(read_fitacf_py, m)?)?; - m.add_function(wrap_pyfunction!(read_snd_py, m)?)?; - m.add_function(wrap_pyfunction!(read_grid_py, m)?)?; - m.add_function(wrap_pyfunction!(read_map_py, m)?)?; - - // Lax read functions - m.add_function(wrap_pyfunction!(read_dmap_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_iqdat_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_rawacf_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_fitacf_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_snd_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_grid_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_map_lax_py, m)?)?; - - // Read functions from byte buffer - m.add_function(wrap_pyfunction!(read_dmap_bytes_py, m)?)?; - m.add_function(wrap_pyfunction!(read_iqdat_bytes_py, m)?)?; - m.add_function(wrap_pyfunction!(read_rawacf_bytes_py, m)?)?; - m.add_function(wrap_pyfunction!(read_fitacf_bytes_py, m)?)?; - m.add_function(wrap_pyfunction!(read_snd_bytes_py, m)?)?; - m.add_function(wrap_pyfunction!(read_grid_bytes_py, m)?)?; - m.add_function(wrap_pyfunction!(read_map_bytes_py, m)?)?; - - // Read select records from byte buffer - m.add_function(wrap_pyfunction!(read_dmap_bytes_by_indices_py, m)?)?; - m.add_function(wrap_pyfunction!(read_iqdat_bytes_by_indices_py, m)?)?; - m.add_function(wrap_pyfunction!(read_rawacf_bytes_by_indices_py, m)?)?; - m.add_function(wrap_pyfunction!(read_fitacf_bytes_by_indices_py, m)?)?; - m.add_function(wrap_pyfunction!(read_snd_bytes_by_indices_py, m)?)?; - m.add_function(wrap_pyfunction!(read_grid_bytes_by_indices_py, m)?)?; - m.add_function(wrap_pyfunction!(read_map_bytes_by_indices_py, m)?)?; - - // Read select records from byte buffer, without raising error - m.add_function(wrap_pyfunction!(read_dmap_bytes_by_indices_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_iqdat_bytes_by_indices_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_rawacf_bytes_by_indices_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_fitacf_bytes_by_indices_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_snd_bytes_by_indices_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_grid_bytes_by_indices_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_map_bytes_by_indices_lax_py, m)?)?; - - // Lax read functions from byte buffer - m.add_function(wrap_pyfunction!(read_dmap_bytes_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_iqdat_bytes_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_rawacf_bytes_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_fitacf_bytes_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_snd_bytes_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_grid_bytes_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_map_bytes_lax_py, m)?)?; - - // Write functions - m.add_function(wrap_pyfunction!(write_dmap_py, m)?)?; - m.add_function(wrap_pyfunction!(write_iqdat_py, m)?)?; - m.add_function(wrap_pyfunction!(write_rawacf_py, m)?)?; - m.add_function(wrap_pyfunction!(write_fitacf_py, m)?)?; - m.add_function(wrap_pyfunction!(write_grid_py, m)?)?; - m.add_function(wrap_pyfunction!(write_map_py, m)?)?; - m.add_function(wrap_pyfunction!(write_snd_py, m)?)?; - - // Convert records to bytes - m.add_function(wrap_pyfunction!(write_dmap_bytes_py, m)?)?; - m.add_function(wrap_pyfunction!(write_iqdat_bytes_py, m)?)?; - m.add_function(wrap_pyfunction!(write_rawacf_bytes_py, m)?)?; - m.add_function(wrap_pyfunction!(write_fitacf_bytes_py, m)?)?; - m.add_function(wrap_pyfunction!(write_snd_bytes_py, m)?)?; - m.add_function(wrap_pyfunction!(write_grid_bytes_py, m)?)?; - m.add_function(wrap_pyfunction!(write_map_bytes_py, m)?)?; - - // Sniff records by index - m.add_function(wrap_pyfunction!(read_dmap_by_indices_py, m)?)?; - m.add_function(wrap_pyfunction!(read_iqdat_by_indices_py, m)?)?; - m.add_function(wrap_pyfunction!(read_rawacf_by_indices_py, m)?)?; - m.add_function(wrap_pyfunction!(read_fitacf_by_indices_py, m)?)?; - m.add_function(wrap_pyfunction!(read_snd_by_indices_py, m)?)?; - m.add_function(wrap_pyfunction!(read_grid_by_indices_py, m)?)?; - m.add_function(wrap_pyfunction!(read_map_by_indices_py, m)?)?; - - // Sniff records by index, but report corrupt records - m.add_function(wrap_pyfunction!(read_dmap_by_indices_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_iqdat_by_indices_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_rawacf_by_indices_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_fitacf_by_indices_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_snd_by_indices_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_grid_by_indices_lax_py, m)?)?; - m.add_function(wrap_pyfunction!(read_map_by_indices_lax_py, m)?)?; - - // Read only the metadata from files - m.add_function(wrap_pyfunction!(read_dmap_metadata_py, m)?)?; - m.add_function(wrap_pyfunction!(read_iqdat_metadata_py, m)?)?; - m.add_function(wrap_pyfunction!(read_rawacf_metadata_py, m)?)?; - m.add_function(wrap_pyfunction!(read_fitacf_metadata_py, m)?)?; - m.add_function(wrap_pyfunction!(read_snd_metadata_py, m)?)?; - m.add_function(wrap_pyfunction!(read_grid_metadata_py, m)?)?; - m.add_function(wrap_pyfunction!(read_map_metadata_py, m)?)?; - - // Read only the metadata of select records from files - m.add_function(wrap_pyfunction!(read_dmap_metadata_by_indices_py, m)?)?; - m.add_function(wrap_pyfunction!(read_iqdat_metadata_by_indices_py, m)?)?; - m.add_function(wrap_pyfunction!(read_rawacf_metadata_by_indices_py, m)?)?; - m.add_function(wrap_pyfunction!(read_fitacf_metadata_by_indices_py, m)?)?; - m.add_function(wrap_pyfunction!(read_snd_metadata_by_indices_py, m)?)?; - m.add_function(wrap_pyfunction!(read_grid_metadata_by_indices_py, m)?)?; - m.add_function(wrap_pyfunction!(read_map_metadata_by_indices_py, m)?)?; - - Ok(()) -} diff --git a/src/parser.rs b/src/parser.rs index c4fa2f1..147e1a9 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -1,8 +1,8 @@ -use std::io::Cursor; -use numpy::ndarray::ArrayD; -use indexmap::IndexMap; -use crate::{DmapError, Record}; use crate::types::{DmapField, DmapScalar, DmapType, DmapVec, Type}; +use crate::{DmapError, Record}; +use indexmap::IndexMap; +use ndarray::ArrayD; +use std::io::Cursor; /// DMAP record header information pub(crate) struct Header { @@ -19,9 +19,7 @@ pub(crate) struct Parser { impl From>> for Parser { fn from(cursor: Cursor>) -> Self { - Self { - cursor, - } + Self { cursor } } } @@ -67,7 +65,8 @@ impl Parser { })?; // adding 8 bytes because code and size are part of the record. - if size as u64 > self.cursor.get_ref().len() as u64 - self.position() + 2 * i32::size() as u64 + if size as u64 + > self.cursor.get_ref().len() as u64 - self.position() + 2 * i32::size() as u64 { return Err(DmapError::InvalidRecord(format!( "Record size {size} at byte {} bigger than remaining buffer {}", @@ -104,7 +103,12 @@ impl Parser { Err(DmapError::InvalidRecord(format!( "Number of scalars {num_scalars} plus vectors {num_vectors} greater than size '{size}'"))) } else { - Ok(Header {size, _code, num_scalars, num_vectors }) + Ok(Header { + size, + _code, + num_scalars, + num_vectors, + }) } } @@ -222,7 +226,8 @@ impl Parser { record_size: i32, ) -> Result<(String, DmapField), DmapError> { let start_position = self.position(); - let (name, data_type, dimensions, total_elements) = self.parse_vector_header(record_size)?; + let (name, data_type, dimensions, total_elements) = + self.parse_vector_header(record_size)?; macro_rules! dmapvec_from_cursor { ($type:ty, $enum_var:path, $dims:ident, $parser:ident, $num_elements:ident, $name:ident) => { @@ -242,14 +247,9 @@ impl Parser { Type::Char => { dmapvec_from_cursor!(i8, DmapVec::Char, dimensions, self, total_elements, name) } - Type::Short => dmapvec_from_cursor!( - i16, - DmapVec::Short, - dimensions, - self, - total_elements, - name - ), + Type::Short => { + dmapvec_from_cursor!(i16, DmapVec::Short, dimensions, self, total_elements, name) + } Type::Int => { dmapvec_from_cursor!(i32, DmapVec::Int, dimensions, self, total_elements, name) } @@ -259,41 +259,21 @@ impl Parser { Type::Uchar => { dmapvec_from_cursor!(u8, DmapVec::Uchar, dimensions, self, total_elements, name) } - Type::Ushort => dmapvec_from_cursor!( - u16, - DmapVec::Ushort, - dimensions, - self, - total_elements, - name - ), + Type::Ushort => { + dmapvec_from_cursor!(u16, DmapVec::Ushort, dimensions, self, total_elements, name) + } Type::Uint => { dmapvec_from_cursor!(u32, DmapVec::Uint, dimensions, self, total_elements, name) } - Type::Ulong => dmapvec_from_cursor!( - u64, - DmapVec::Ulong, - dimensions, - self, - total_elements, - name - ), - Type::Float => dmapvec_from_cursor!( - f32, - DmapVec::Float, - dimensions, - self, - total_elements, - name - ), - Type::Double => dmapvec_from_cursor!( - f64, - DmapVec::Double, - dimensions, - self, - total_elements, - name - ), + Type::Ulong => { + dmapvec_from_cursor!(u64, DmapVec::Ulong, dimensions, self, total_elements, name) + } + Type::Float => { + dmapvec_from_cursor!(f32, DmapVec::Float, dimensions, self, total_elements, name) + } + Type::Double => { + dmapvec_from_cursor!(f64, DmapVec::Double, dimensions, self, total_elements, name) + } Type::String => { return Err(DmapError::InvalidVector(format!( "Invalid type {data_type} for DMAP vector {name}" @@ -382,13 +362,15 @@ impl Parser { ))); } - self.set_position(0); // reset the cursor to the start + self.set_position(0); // reset the cursor to the start T::new(&mut fields) } /// Reads a record from `self`, only keeping the metadata fields. - pub(crate) fn parse_metadata<'a, T: Record<'a>>(&mut self) -> Result, DmapError> + pub(crate) fn parse_metadata<'a, T: Record<'a>>( + &mut self, + ) -> Result, DmapError> where Self: Sized, { @@ -427,12 +409,11 @@ impl Parser { #[cfg(test)] mod tests { - use numpy::array; use super::*; + use ndarray::array; #[test] fn test_read_vec() { - let bytes: Vec = vec![1, 0, 1, 0]; let mut parser = Parser::new(bytes.clone()); let data = parser.read_vector::(4); @@ -672,4 +653,4 @@ mod tests { Ok(()) } -} \ No newline at end of file +} diff --git a/src/python_bindings.rs b/src/python_bindings.rs new file mode 100644 index 0000000..062d4a9 --- /dev/null +++ b/src/python_bindings.rs @@ -0,0 +1,462 @@ +use crate::compression; +use crate::formats::dmap::DmapRecord; +use crate::formats::fitacf::FitacfRecord; +use crate::formats::grid::GridRecord; +use crate::formats::iqdat::IqdatRecord; +use crate::formats::map::MapRecord; +use crate::formats::rawacf::RawacfRecord; +use crate::formats::snd::SndRecord; +use crate::record::Record; +use crate::types::DmapField; +use indexmap::IndexMap; +use paste::paste; +use pyo3::prelude::*; +use pyo3::types::PyBytes; +use pyo3::{Bound, PyAny, PyResult, Python}; +use std::path::PathBuf; + +/// Creates functions for reading DMAP files for the Python API. +/// +/// Generates functions for: +/// +/// * `read_[name]` - reads a file, raising an error on a corrupted file +/// * `read_[name]_lax` - reads a file, returning the records and the byte where corruption starts, if corrupted. +/// * `read_[name]_bytes` - reads from bytes, similar to `read_[name]` +/// * `read_[name]_bytes_lax` - reads from bytes, similar to `read_[name]_lax` +/// * `read_[name]_by_indices` - reads specific records from a file by index +/// * `read_[name]_by_indices_lax` - reads specific records from a file by index, returning the byte offset where corruption begins, if applicable +/// * `read_[name]_by_indices_bytes` - reads specific records from a byte buffer by index +/// * `read_[name]_by_indices_bytes_lax` - reads specific records from a byte buffer by index, returning the byte offset where corruption begins, if applicable +/// * `read_[name]_metadata` - reads only metadata fields from records in a file +/// * `read_[name]_metadata_by_indices` - reads only metadata fields from specific records in a file by index +/// +/// where `[name]` is one of the supported DMAP record types. +macro_rules! read_py { + ( + $name:ident, + $py_name:literal, + $lax_name:literal, + $bytes_name:literal, + $lax_bytes_name:literal, + $by_indices_name:literal, + $by_indices_name_lax:literal, + $bytes_by_indices_name:literal, + $bytes_by_indices_name_lax:literal, + $metadata_name:literal, + $metadata_by_indices_name:literal + ) => { + paste! { + #[doc = "Reads a `" $name:upper "` file, returning a list of dictionaries containing the fields." ] + #[pyfunction] + #[pyo3(name = $py_name)] + #[pyo3(text_signature = "(infile: str, /)")] + fn [< read_ $name _py >](infile: PathBuf) -> PyResult>> { + Ok([< $name:camel Record >]::read_file(&infile) + .map_err(PyErr::from)? + .into_iter() + .map(|rec| rec.inner()) + .collect() + ) + } + + #[doc = "Reads a `" $name:upper "` file, returning a tuple of" ] + #[doc = "(list of dictionaries containing the fields, byte where first corrupted record starts). "] + #[pyfunction] + #[pyo3(name = $lax_name)] + #[pyo3(text_signature = "(infile: str, /)")] + fn [< read_ $name _lax_py >]( + infile: PathBuf, + ) -> PyResult<(Vec>, Option)> { + let result = [< $name:camel Record >]::read_file_lax(&infile).map_err(PyErr::from)?; + Ok(( + result.0.into_iter().map(|rec| rec.inner()).collect(), + result.1, + )) + } + + #[doc = "Read in `" $name:upper "` records from bytes, returning `List[Dict]` of the records." ] + #[pyfunction] + #[pyo3(name = $bytes_name)] + #[pyo3(text_signature = "(buf: bytes, /)")] + fn [< read_ $name _bytes_py >](bytes: &[u8]) -> PyResult>> { + Ok([< $name:camel Record >]::read_records(bytes)? + .into_iter() + .map(|rec| rec.inner()) + .collect() + ) + } + + #[doc = "Reads a `" $name:upper "` file, returning a tuple of" ] + #[doc = "(list of dictionaries containing the fields, byte where first corrupted record starts). "] + #[pyfunction] + #[pyo3(name = $lax_bytes_name)] + #[pyo3(text_signature = "(buf: bytes, /)")] + fn [< read_ $name _bytes_lax_py >]( + bytes: &[u8], + ) -> PyResult<(Vec>, Option)> { + let result = [< $name:camel Record >]::read_records_lax(bytes).map_err(PyErr::from)?; + Ok(( + result.0.into_iter().map(|rec| rec.inner()).collect(), + result.1, + )) + } + + #[doc = "Reads a `" $name:upper "` file, returning the `nth` record(s)." ] + #[pyfunction] + #[pyo3(name = $by_indices_name)] + #[pyo3(text_signature = "(infile: str, indices: tuple[int], /)")] + fn [< read_ $name _by_indices_py >](infile: PathBuf, indices: Vec) -> PyResult>> { + Ok([< $name:camel Record >]::read_file_by_indices(&infile, &indices) + .map_err(PyErr::from)? + .into_iter() + .map(|rec| rec.inner()) + .collect() + ) + } + + #[doc = "Reads a `" $name:upper "` file, returning the `nth` record(s), and the byte index where corruption starts, if applicable" ] + #[pyfunction] + #[pyo3(name = $by_indices_name_lax)] + #[pyo3(text_signature = "(infile: str, indices: tuple[int], /)")] + fn [< read_ $name _by_indices_lax_py >]( + infile: PathBuf, indices: Vec + ) -> PyResult<(Vec>, Option)> { + let result = [< $name:camel Record >]::read_file_by_indices_lax(infile, &indices).map_err(PyErr::from)?; + Ok(( + result.0.into_iter().map(|rec| rec.inner()).collect(), + result.1, + )) + } + + #[doc = "Reads a `" $name:upper "` buffer, returning the `nth` record(s)." ] + #[pyfunction] + #[pyo3(name = $bytes_by_indices_name)] + #[pyo3(text_signature = "(buf: bytes, indices: tuple[int], /)")] + fn [< read_ $name _bytes_by_indices_py >](buf: &[u8], indices: Vec) -> PyResult>> { + Ok([< $name:camel Record >]::read_nth_records(buf, &indices) + .map_err(PyErr::from)? + .into_iter() + .map(|rec| rec.inner()) + .collect() + ) + } + + #[doc = "Reads a `" $name:upper "` buffer, returning the `nth` record(s) and the byte index where record corruption starts, if applicable." ] + #[pyfunction] + #[pyo3(name = $bytes_by_indices_name_lax)] + #[pyo3(text_signature = "(buf: bytes, indices: tuple[int], /)")] + fn [< read_ $name _bytes_by_indices_lax_py >]( + buf: &[u8], indices: Vec + ) -> PyResult<(Vec>, Option)> { + let result = [< $name:camel Record >]::read_nth_records_lax(buf, &indices).map_err(PyErr::from)?; + Ok(( + result.0.into_iter().map(|rec| rec.inner()).collect(), + result.1, + )) + } + + #[doc = "Reads a `" $name:upper "` file, returning a list of dictionaries containing the only the metadata fields." ] + #[pyfunction] + #[pyo3(name = $metadata_name)] + #[pyo3(text_signature = "(infile: str, /)")] + fn [< read_ $name _metadata_py >](infile: PathBuf) -> PyResult>> { + Ok([< $name:camel Record >]::read_file_metadata(&infile) + .map_err(PyErr::from)? + ) + } + + #[doc = "Reads a `" $name:upper "` file, returning the `nth` records' metadata fields." ] + #[pyfunction] + #[pyo3(name = $metadata_by_indices_name)] + #[pyo3(text_signature = "(infile: str, indices: tuple[int], /)")] + fn [< read_ $name _metadata_by_indices_py >](infile: PathBuf, indices: Vec) -> PyResult>> { + Ok([< $name:camel Record >]::read_file_metadata_by_indices(&infile, &indices) + .map_err(PyErr::from)? + ) + } + } + } +} + +read_py!( + iqdat, + "read_iqdat", + "read_iqdat_lax", + "read_iqdat_bytes", + "read_iqdat_bytes_lax", + "read_iqdat_by_indices", + "read_iqdat_by_indices_lax", + "read_iqdat_by_indices_bytes", + "read_iqdat_by_indices_bytes_lax", + "read_iqdat_metadata", + "read_iqdat_metadata_by_indices" +); +read_py!( + rawacf, + "read_rawacf", + "read_rawacf_lax", + "read_rawacf_bytes", + "read_rawacf_bytes_lax", + "read_rawacf_by_indices", + "read_rawacf_by_indices_lax", + "read_rawacf_by_indices_bytes", + "read_rawacf_by_indices_bytes_lax", + "read_rawacf_metadata", + "read_rawacf_metadata_by_indices" +); +read_py!( + fitacf, + "read_fitacf", + "read_fitacf_lax", + "read_fitacf_bytes", + "read_fitacf_bytes_lax", + "read_fitacf_by_indices", + "read_fitacf_by_indices_lax", + "read_fitacf_by_indices_bytes", + "read_fitacf_by_indices_bytes_lax", + "read_fitacf_metadata", + "read_fitacf_metadata_by_indices" +); +read_py!( + grid, + "read_grid", + "read_grid_lax", + "read_grid_bytes", + "read_grid_bytes_lax", + "read_grid_by_indices", + "read_grid_by_indices_lax", + "read_grid_by_indices_bytes", + "read_grid_by_indices_bytes_lax", + "read_grid_metadata", + "read_grid_metadata_by_indices" +); +read_py!( + map, + "read_map", + "read_map_lax", + "read_map_bytes", + "read_map_bytes_lax", + "read_map_by_indices", + "read_map_by_indices_lax", + "read_map_by_indices_bytes", + "read_map_by_indices_bytes_lax", + "read_map_metadata", + "read_map_metadata_by_indices" +); +read_py!( + snd, + "read_snd", + "read_snd_lax", + "read_snd_bytes", + "read_snd_bytes_lax", + "read_snd_by_indices", + "read_snd_by_indices_lax", + "read_snd_by_indices_bytes", + "read_snd_by_indices_bytes_lax", + "read_snd_metadata", + "read_snd_metadata_by_indices" +); +read_py!( + dmap, + "read_dmap", + "read_dmap_lax", + "read_dmap_bytes", + "read_dmap_bytes_lax", + "read_dmap_by_indices", + "read_dmap_by_indices_lax", + "read_dmap_by_indices_bytes", + "read_dmap_by_indices_bytes_lax", + "read_dmap_metadata", + "read_dmap_metadata_by_indices" +); + +/// Checks that a list of dictionaries contains DMAP records, then appends to outfile. +/// +/// **NOTE:** No type checking is done, so the fields may not be written as the expected +/// DMAP type, e.g. `stid` might be written one byte instead of two as this function +/// does not know that typically `stid` is two bytes. +#[pyfunction] +#[pyo3(name = "write_dmap")] +#[pyo3(signature = (recs, outfile, /, bz2))] +#[pyo3(text_signature = "(recs: list[dict], outfile: str, /, bz2: bool = False)")] +fn write_dmap_py( + recs: Vec>, + outfile: PathBuf, + bz2: bool, +) -> PyResult<()> { + DmapRecord::try_write_to_file(recs, &outfile, bz2).map_err(PyErr::from) +} + +/// Checks that a list of dictionaries contains valid DMAP records, then converts them to bytes. +/// Returns a `bytes` object containing the serialized records. +/// +/// **NOTE:** No type checking is done, so the fields may not be written as the expected +/// DMAP type, e.g. `stid` might be written one byte instead of two as this function +/// does not know that typically `stid` is two bytes. +#[pyfunction] +#[pyo3(name = "write_dmap_bytes")] +#[pyo3(signature = (recs, /, bz2))] +#[pyo3(text_signature = "(recs: list[dict], /, bz2: bool = False)")] +fn write_dmap_bytes_py( + py: Python, + recs: Vec>, + bz2: bool, +) -> PyResult> { + let mut bytes = DmapRecord::try_into_bytes(recs).map_err(PyErr::from)?; + if bz2 { + bytes = compression::compress_bz2(&bytes).map_err(PyErr::from)?; + } + Ok(PyBytes::new(py, &bytes).into()) +} + +/// Generates functions exposed to the Python API for writing specific file types. +macro_rules! write_py { + ($name:ident, $fn_name:literal, $bytes_name:literal) => { + paste! { + #[doc = "Checks that a list of dictionaries contains valid `" $name:upper "` records, then appends to outfile." ] + #[pyfunction] + #[pyo3(name = $fn_name)] + #[pyo3(signature = (recs, outfile, /, bz2))] + #[pyo3(text_signature = "(recs: list[dict], outfile: str, /, bz2: bool = False)")] + fn [< write_ $name _py >](recs: Vec>, outfile: PathBuf, bz2: bool) -> PyResult<()> { + [< $name:camel Record >]::try_write_to_file(recs, &outfile, bz2).map_err(PyErr::from) + } + + #[doc = "Checks that a list of dictionaries contains valid `" $name:upper "` records, then converts them to bytes." ] + #[doc = "Returns a `bytes` object containing the serialized records." ] + #[pyfunction] + #[pyo3(name = $bytes_name)] + #[pyo3(signature = (recs, /, bz2))] + #[pyo3(text_signature = "(recs: list[dict], /, bz2: bool = False)")] + fn [< write_ $name _bytes_py >](py: Python, recs: Vec>, bz2: bool) -> PyResult> { + let mut bytes = [< $name:camel Record >]::try_into_bytes(recs).map_err(PyErr::from)?; + if bz2 { + bytes = compression::compress_bz2(&bytes).map_err(PyErr::from)?; + } + Ok(PyBytes::new(py, &bytes).into()) + } + } + } +} + +// **NOTE** dmap type not included in this list, since it has a more descriptive docstring. +write_py!(iqdat, "write_iqdat", "write_iqdat_bytes"); +write_py!(rawacf, "write_rawacf", "write_rawacf_bytes"); +write_py!(fitacf, "write_fitacf", "write_fitacf_bytes"); +write_py!(grid, "write_grid", "write_grid_bytes"); +write_py!(map, "write_map", "write_map_bytes"); +write_py!(snd, "write_snd", "write_snd_bytes"); + +/// Functions for SuperDARN DMAP file format I/O. +#[pymodule] +fn dmap_rs(m: &Bound<'_, PyModule>) -> PyResult<()> { + // Strict read functions + m.add_function(wrap_pyfunction!(read_dmap_py, m)?)?; + m.add_function(wrap_pyfunction!(read_iqdat_py, m)?)?; + m.add_function(wrap_pyfunction!(read_rawacf_py, m)?)?; + m.add_function(wrap_pyfunction!(read_fitacf_py, m)?)?; + m.add_function(wrap_pyfunction!(read_snd_py, m)?)?; + m.add_function(wrap_pyfunction!(read_grid_py, m)?)?; + m.add_function(wrap_pyfunction!(read_map_py, m)?)?; + + // Lax read functions + m.add_function(wrap_pyfunction!(read_dmap_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_iqdat_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_rawacf_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_fitacf_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_snd_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_grid_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_map_lax_py, m)?)?; + + // Read functions from byte buffer + m.add_function(wrap_pyfunction!(read_dmap_bytes_py, m)?)?; + m.add_function(wrap_pyfunction!(read_iqdat_bytes_py, m)?)?; + m.add_function(wrap_pyfunction!(read_rawacf_bytes_py, m)?)?; + m.add_function(wrap_pyfunction!(read_fitacf_bytes_py, m)?)?; + m.add_function(wrap_pyfunction!(read_snd_bytes_py, m)?)?; + m.add_function(wrap_pyfunction!(read_grid_bytes_py, m)?)?; + m.add_function(wrap_pyfunction!(read_map_bytes_py, m)?)?; + + // Read select records from byte buffer + m.add_function(wrap_pyfunction!(read_dmap_bytes_by_indices_py, m)?)?; + m.add_function(wrap_pyfunction!(read_iqdat_bytes_by_indices_py, m)?)?; + m.add_function(wrap_pyfunction!(read_rawacf_bytes_by_indices_py, m)?)?; + m.add_function(wrap_pyfunction!(read_fitacf_bytes_by_indices_py, m)?)?; + m.add_function(wrap_pyfunction!(read_snd_bytes_by_indices_py, m)?)?; + m.add_function(wrap_pyfunction!(read_grid_bytes_by_indices_py, m)?)?; + m.add_function(wrap_pyfunction!(read_map_bytes_by_indices_py, m)?)?; + + // Read select records from byte buffer, without raising error + m.add_function(wrap_pyfunction!(read_dmap_bytes_by_indices_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_iqdat_bytes_by_indices_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_rawacf_bytes_by_indices_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_fitacf_bytes_by_indices_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_snd_bytes_by_indices_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_grid_bytes_by_indices_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_map_bytes_by_indices_lax_py, m)?)?; + + // Lax read functions from byte buffer + m.add_function(wrap_pyfunction!(read_dmap_bytes_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_iqdat_bytes_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_rawacf_bytes_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_fitacf_bytes_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_snd_bytes_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_grid_bytes_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_map_bytes_lax_py, m)?)?; + + // Write functions + m.add_function(wrap_pyfunction!(write_dmap_py, m)?)?; + m.add_function(wrap_pyfunction!(write_iqdat_py, m)?)?; + m.add_function(wrap_pyfunction!(write_rawacf_py, m)?)?; + m.add_function(wrap_pyfunction!(write_fitacf_py, m)?)?; + m.add_function(wrap_pyfunction!(write_grid_py, m)?)?; + m.add_function(wrap_pyfunction!(write_map_py, m)?)?; + m.add_function(wrap_pyfunction!(write_snd_py, m)?)?; + + // Convert records to bytes + m.add_function(wrap_pyfunction!(write_dmap_bytes_py, m)?)?; + m.add_function(wrap_pyfunction!(write_iqdat_bytes_py, m)?)?; + m.add_function(wrap_pyfunction!(write_rawacf_bytes_py, m)?)?; + m.add_function(wrap_pyfunction!(write_fitacf_bytes_py, m)?)?; + m.add_function(wrap_pyfunction!(write_snd_bytes_py, m)?)?; + m.add_function(wrap_pyfunction!(write_grid_bytes_py, m)?)?; + m.add_function(wrap_pyfunction!(write_map_bytes_py, m)?)?; + + // Read records by index + m.add_function(wrap_pyfunction!(read_dmap_by_indices_py, m)?)?; + m.add_function(wrap_pyfunction!(read_iqdat_by_indices_py, m)?)?; + m.add_function(wrap_pyfunction!(read_rawacf_by_indices_py, m)?)?; + m.add_function(wrap_pyfunction!(read_fitacf_by_indices_py, m)?)?; + m.add_function(wrap_pyfunction!(read_snd_by_indices_py, m)?)?; + m.add_function(wrap_pyfunction!(read_grid_by_indices_py, m)?)?; + m.add_function(wrap_pyfunction!(read_map_by_indices_py, m)?)?; + + // Read records by index, but report corrupt records + m.add_function(wrap_pyfunction!(read_dmap_by_indices_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_iqdat_by_indices_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_rawacf_by_indices_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_fitacf_by_indices_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_snd_by_indices_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_grid_by_indices_lax_py, m)?)?; + m.add_function(wrap_pyfunction!(read_map_by_indices_lax_py, m)?)?; + + // Read only the metadata from files + m.add_function(wrap_pyfunction!(read_dmap_metadata_py, m)?)?; + m.add_function(wrap_pyfunction!(read_iqdat_metadata_py, m)?)?; + m.add_function(wrap_pyfunction!(read_rawacf_metadata_py, m)?)?; + m.add_function(wrap_pyfunction!(read_fitacf_metadata_py, m)?)?; + m.add_function(wrap_pyfunction!(read_snd_metadata_py, m)?)?; + m.add_function(wrap_pyfunction!(read_grid_metadata_py, m)?)?; + m.add_function(wrap_pyfunction!(read_map_metadata_py, m)?)?; + + // Read only the metadata of select records from files + m.add_function(wrap_pyfunction!(read_dmap_metadata_by_indices_py, m)?)?; + m.add_function(wrap_pyfunction!(read_iqdat_metadata_by_indices_py, m)?)?; + m.add_function(wrap_pyfunction!(read_rawacf_metadata_by_indices_py, m)?)?; + m.add_function(wrap_pyfunction!(read_fitacf_metadata_by_indices_py, m)?)?; + m.add_function(wrap_pyfunction!(read_snd_metadata_by_indices_py, m)?)?; + m.add_function(wrap_pyfunction!(read_grid_metadata_by_indices_py, m)?)?; + m.add_function(wrap_pyfunction!(read_map_metadata_by_indices_py, m)?)?; + + Ok(()) +} diff --git a/src/python_types.rs b/src/python_types.rs new file mode 100644 index 0000000..2d8cb94 --- /dev/null +++ b/src/python_types.rs @@ -0,0 +1,132 @@ +use crate::types::{DmapField, DmapScalar, DmapVec}; + +use numpy::{PyArray, PyArrayMethods}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::{Bound, FromPyObject, PyAny, PyResult, Python}; + +impl<'py> IntoPyObject<'py> for DmapScalar { + type Target = PyAny; + type Output = Bound<'py, Self::Target>; + type Error = std::convert::Infallible; + + fn into_pyobject(self, py: Python<'py>) -> Result { + Ok(match self { + Self::Char(x) => x.into_pyobject(py)?.into_any(), + Self::Short(x) => x.into_pyobject(py)?.into_any(), + Self::Int(x) => x.into_pyobject(py)?.into_any(), + Self::Long(x) => x.into_pyobject(py)?.into_any(), + Self::Uchar(x) => x.into_pyobject(py)?.into_any(), + Self::Ushort(x) => x.into_pyobject(py)?.into_any(), + Self::Uint(x) => x.into_pyobject(py)?.into_any(), + Self::Ulong(x) => x.into_pyobject(py)?.into_any(), + Self::Float(x) => x.into_pyobject(py)?.into_any(), + Self::Double(x) => x.into_pyobject(py)?.into_any(), + Self::String(x) => x.into_pyobject(py)?.into_any(), + }) + } +} + +impl<'py> FromPyObject<'py> for DmapScalar { + fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { + if let Ok(x) = ob.extract::() { + Ok(DmapScalar::String(x)) + } else if let Ok(x) = ob.extract::() { + Ok(DmapScalar::Char(x)) + } else if let Ok(x) = ob.extract::() { + Ok(DmapScalar::Short(x)) + } else if let Ok(x) = ob.extract::() { + Ok(DmapScalar::Int(x)) + } else if let Ok(x) = ob.extract::() { + Ok(DmapScalar::Long(x)) + } else if let Ok(x) = ob.extract::() { + Ok(DmapScalar::Uchar(x)) + } else if let Ok(x) = ob.extract::() { + Ok(DmapScalar::Ushort(x)) + } else if let Ok(x) = ob.extract::() { + Ok(DmapScalar::Uint(x)) + } else if let Ok(x) = ob.extract::() { + Ok(DmapScalar::Ulong(x)) + } else if let Ok(x) = ob.extract::() { + Ok(DmapScalar::Float(x)) + } else if let Ok(x) = ob.extract::() { + Ok(DmapScalar::Double(x)) + } else { + Err(PyValueError::new_err("Could not extract scalar")) + } + } +} + +impl<'py> IntoPyObject<'py> for DmapVec { + type Target = PyAny; + type Output = Bound<'py, Self::Target>; + type Error = std::convert::Infallible; + + fn into_pyobject(self, py: Python<'py>) -> std::result::Result { + Ok(match self { + DmapVec::Char(x) => PyArray::from_owned_array(py, x).into_any(), + DmapVec::Short(x) => PyArray::from_owned_array(py, x).into_any(), + DmapVec::Int(x) => PyArray::from_owned_array(py, x).into_any(), + DmapVec::Long(x) => PyArray::from_owned_array(py, x).into_any(), + DmapVec::Uchar(x) => PyArray::from_owned_array(py, x).into_any(), + DmapVec::Ushort(x) => PyArray::from_owned_array(py, x).into_any(), + DmapVec::Uint(x) => PyArray::from_owned_array(py, x).into_any(), + DmapVec::Ulong(x) => PyArray::from_owned_array(py, x).into_any(), + DmapVec::Float(x) => PyArray::from_owned_array(py, x).into_any(), + DmapVec::Double(x) => PyArray::from_owned_array(py, x).into_any(), + }) + } +} + +impl<'py> FromPyObject<'py> for DmapVec { + fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { + if let Ok(x) = ob.downcast::>() { + Ok(DmapVec::Uchar(x.to_owned_array())) + } else if let Ok(x) = ob.downcast::>() { + Ok(DmapVec::Ushort(x.to_owned_array())) + } else if let Ok(x) = ob.downcast::>() { + Ok(DmapVec::Uint(x.to_owned_array())) + } else if let Ok(x) = ob.downcast::>() { + Ok(DmapVec::Ulong(x.to_owned_array())) + } else if let Ok(x) = ob.downcast::>() { + Ok(DmapVec::Char(x.to_owned_array())) + } else if let Ok(x) = ob.downcast::>() { + Ok(DmapVec::Short(x.to_owned_array())) + } else if let Ok(x) = ob.downcast::>() { + Ok(DmapVec::Int(x.to_owned_array())) + } else if let Ok(x) = ob.downcast::>() { + Ok(DmapVec::Long(x.to_owned_array())) + } else if let Ok(x) = ob.downcast::>() { + Ok(DmapVec::Float(x.to_owned_array())) + } else if let Ok(x) = ob.downcast::>() { + Ok(DmapVec::Double(x.to_owned_array())) + } else { + Err(PyValueError::new_err("Could not extract vector")) + } + } +} + +impl<'py> IntoPyObject<'py> for DmapField { + type Target = PyAny; + type Output = Bound<'py, Self::Target>; + type Error = std::convert::Infallible; + + fn into_pyobject(self, py: Python<'py>) -> Result { + Ok(match self { + DmapField::Scalar(x) => x.into_pyobject(py)?.into_any(), + DmapField::Vector(x) => x.into_pyobject(py)?.into_any(), + }) + } +} + +impl<'py> FromPyObject<'py> for DmapField { + fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { + if let Ok(x) = ob.extract::() { + Ok(DmapField::Vector(x)) + } else if let Ok(x) = ob.extract::() { + Ok(DmapField::Scalar(x)) + } else { + Err(PyValueError::new_err("Could not extract field")) + } + } +} diff --git a/src/record.rs b/src/record.rs index 361552f..d333574 100644 --- a/src/record.rs +++ b/src/record.rs @@ -1,10 +1,10 @@ //! Defines the [`Record`] trait, which contains the shared behaviour that all DMAP records must have. +use crate::convenience::split_results; use crate::error::DmapError; -use crate::io; -use crate::types::{DmapField, DmapType, DmapVec, Fields}; +use crate::io::{self}; use crate::io::{create_stream, slice_stream_lax, split_into_slices}; -use crate::convenience::split_results; +use crate::types::{DmapField, DmapType, DmapVec, Fields}; use indexmap::IndexMap; use itertools::izip; use rayon::iter::Either; @@ -60,8 +60,9 @@ pub trait Record<'a>: *idx as usize } }; - records_read.push(slices[i].parse_record::() - .map_err(|e| DmapError::InvalidRecord(format!("Record {idx}: {}", e.to_string())))?); + records_read.push(slices[i].parse_record::().map_err(|e| { + DmapError::InvalidRecord(format!("Record {idx}: {}", e.to_string())) + })?); } Ok(records_read) } @@ -70,7 +71,10 @@ pub trait Record<'a>: /// /// Returns a 2-tuple, where the first entry is the good records from the front of the buffer, /// and the second entry is the byte where the first corrupted record starts, if applicable. - fn read_nth_records_lax(mut dmap_data: impl Read, indices: &[i32]) -> Result<(Vec, Option), DmapError> + fn read_nth_records_lax( + mut dmap_data: impl Read, + indices: &[i32], + ) -> Result<(Vec, Option), DmapError> where Self: Sized, Self: Send, @@ -155,7 +159,10 @@ pub trait Record<'a>: /// Reads metadata of the nth records from `dmap_data` and parses into `Vec>`. /// /// Returns `DmapError` if `dmap_data` cannot be read or contains invalid data. - fn read_metadata_by_indices(dmap_data: impl Read, indices: &[i32]) -> Result>, DmapError> + fn read_metadata_by_indices( + dmap_data: impl Read, + indices: &[i32], + ) -> Result>, DmapError> where Self: Sized, Self: Send, @@ -175,8 +182,9 @@ pub trait Record<'a>: *idx as usize } }; - dmap_results.push(slices[i].parse_metadata::() - .map_err(|e| DmapError::InvalidRecord(format!("Record {idx}: {}", e.to_string())))?); + dmap_results.push(slices[i].parse_metadata::().map_err(|e| { + DmapError::InvalidRecord(format!("Record {idx}: {}", e.to_string())) + })?); } Ok(dmap_results) } @@ -238,7 +246,10 @@ pub trait Record<'a>: } /// Reads the `nth` record(s) of a DMAP file of type `Self`. - fn read_file_by_indices>(infile: P, indices: &[i32]) -> Result, DmapError> + fn read_file_by_indices>( + infile: P, + indices: &[i32], + ) -> Result, DmapError> where Self: Sized, Self: Send, @@ -251,7 +262,10 @@ pub trait Record<'a>: /// /// Does not fail on corrupted records; rather, returns `(recs, Option)` where /// the second value is the byte where record corruption begins, if applicable. - fn read_file_by_indices_lax>(infile: P, indices: &[i32]) -> Result<(Vec, Option), DmapError> + fn read_file_by_indices_lax>( + infile: P, + indices: &[i32], + ) -> Result<(Vec, Option), DmapError> where Self: Sized, Self: Send, @@ -273,7 +287,10 @@ pub trait Record<'a>: } /// Reads the `nth` records' metadata of a DMAP file of type `Self`. - fn read_file_metadata_by_indices>(infile: P, indices: &[i32]) -> Result>, DmapError> + fn read_file_metadata_by_indices>( + infile: P, + indices: &[i32], + ) -> Result>, DmapError> where Self: Sized, Self: Send, @@ -281,7 +298,7 @@ pub trait Record<'a>: let file = File::open(infile)?; Self::read_metadata_by_indices(file, indices) } - + /// Checks the validity of an `IndexMap` as a representation of a DMAP record. /// /// Validity checks include ensuring that no unfamiliar entries exist, that all required @@ -729,8 +746,9 @@ pub trait Record<'a>: /// Writes a collection of `Record`s to `outfile`. /// - /// Prefer using the specific functions, e.g. `write_dmap`, `write_rawacf`, etc. for their - /// specific field checks. + /// See also [`Record::try_write_to_file`] for writing generic + /// `IndexMap` records that must first be validated + /// against `Self`. fn write_to_file>( recs: &Vec, outfile: P, @@ -740,6 +758,27 @@ pub trait Record<'a>: io::bytes_to_file(bytes, outfile, bz2)?; Ok(()) } + + /// Attempts to convert a collection of generic DMAP records to `Self` and + /// write them to `outfile`. + /// + /// This is useful when the records are represented as + /// `IndexMap` values and need to be validated against the + /// requirements of a specific DMAP format before being written. + /// + /// # Errors + /// + /// Returns an error if any record cannot be converted to `Self`, or if + /// writing to `outfile` fails. + fn try_write_to_file>( + recs: Vec>, + outfile: P, + bz2: bool, + ) -> Result<(), DmapError> { + let bytes = Self::try_into_bytes(recs)?; + io::bytes_to_file(bytes, outfile, bz2)?; + Ok(()) + } } macro_rules! create_record_type { diff --git a/src/types.rs b/src/types.rs index 47f260a..48a6776 100644 --- a/src/types.rs +++ b/src/types.rs @@ -1,13 +1,8 @@ //! Low-level data types within DMAP records. use crate::error::DmapError; use indexmap::IndexMap; -use numpy::array::PyArray; -use numpy::ndarray::ArrayD; -use numpy::PyArrayMethods; +use ndarray::ArrayD; use paste::paste; -use pyo3::exceptions::PyValueError; -use pyo3::prelude::*; -use pyo3::{Bound, FromPyObject, PyAny, PyResult, Python}; use std::cmp::PartialEq; use std::fmt::{Display, Formatter}; use zerocopy::{AsBytes, ByteOrder, FromBytes, LittleEndian}; @@ -121,7 +116,8 @@ impl Type { } /// A scalar field in a DMAP record. -#[derive(Debug, Clone, PartialEq, FromPyObject, IntoPyObject)] +//#[derive(Debug, Clone, PartialEq, FromPyObject, IntoPyObject)] +#[derive(Debug, Clone, PartialEq)] #[repr(C)] pub enum DmapScalar { Char(i8), @@ -209,23 +205,6 @@ impl Display for DmapScalar { } } } -// impl IntoPy for DmapScalar { -// fn into_py(self, py: Python<'_>) -> PyObject { -// match self { -// Self::Char(x) => x.into_py(py), -// Self::Short(x) => x.into_py(py), -// Self::Int(x) => x.into_py(py), -// Self::Long(x) => x.into_py(py), -// Self::Uchar(x) => x.into_py(py), -// Self::Ushort(x) => x.into_py(py), -// Self::Uint(x) => x.into_py(py), -// Self::Ulong(x) => x.into_py(py), -// Self::Float(x) => x.into_py(py), -// Self::Double(x) => x.into_py(py), -// Self::String(x) => x.into_py(py), -// } -// } -// } macro_rules! vec_to_bytes { ($bytes:ident, $x:ident) => {{ @@ -292,7 +271,7 @@ impl DmapVec { /// Gets the dimensions of the vector, in row-major order. /// ## Example /// ``` - /// use numpy::ndarray::array; + /// use ndarray::array; /// use dmap::types::DmapVec; /// /// let arr = DmapVec::Char(array![0, 1, 2, 3, 4].into_dyn()); @@ -317,53 +296,6 @@ impl DmapVec { } } } -impl<'py> IntoPyObject<'py> for DmapVec { - type Target = PyAny; - type Output = Bound<'py, Self::Target>; - type Error = std::convert::Infallible; - - fn into_pyobject(self, py: Python<'py>) -> std::result::Result { - Ok(match self { - DmapVec::Char(x) => PyArray::from_owned_array(py, x).into_any(), - DmapVec::Short(x) => PyArray::from_owned_array(py, x).into_any(), - DmapVec::Int(x) => PyArray::from_owned_array(py, x).into_any(), - DmapVec::Long(x) => PyArray::from_owned_array(py, x).into_any(), - DmapVec::Uchar(x) => PyArray::from_owned_array(py, x).into_any(), - DmapVec::Ushort(x) => PyArray::from_owned_array(py, x).into_any(), - DmapVec::Uint(x) => PyArray::from_owned_array(py, x).into_any(), - DmapVec::Ulong(x) => PyArray::from_owned_array(py, x).into_any(), - DmapVec::Float(x) => PyArray::from_owned_array(py, x).into_any(), - DmapVec::Double(x) => PyArray::from_owned_array(py, x).into_any(), - }) - } -} -impl<'py> FromPyObject<'py> for DmapVec { - fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { - if let Ok(x) = ob.downcast::>() { - Ok(DmapVec::Uchar(x.to_owned_array())) - } else if let Ok(x) = ob.downcast::>() { - Ok(DmapVec::Ushort(x.to_owned_array())) - } else if let Ok(x) = ob.downcast::>() { - Ok(DmapVec::Uint(x.to_owned_array())) - } else if let Ok(x) = ob.downcast::>() { - Ok(DmapVec::Ulong(x.to_owned_array())) - } else if let Ok(x) = ob.downcast::>() { - Ok(DmapVec::Char(x.to_owned_array())) - } else if let Ok(x) = ob.downcast::>() { - Ok(DmapVec::Short(x.to_owned_array())) - } else if let Ok(x) = ob.downcast::>() { - Ok(DmapVec::Int(x.to_owned_array())) - } else if let Ok(x) = ob.downcast::>() { - Ok(DmapVec::Long(x.to_owned_array())) - } else if let Ok(x) = ob.downcast::>() { - Ok(DmapVec::Float(x.to_owned_array())) - } else if let Ok(x) = ob.downcast::>() { - Ok(DmapVec::Double(x.to_owned_array())) - } else { - Err(PyValueError::new_err("Could not extract vector")) - } - } -} /// Generates trait implementations for infallible conversion into [`DmapVec`] and fallible conversion /// back. @@ -432,7 +364,7 @@ vec_impls!(ArrayD, DmapVec::Double); /// /// This is the type that is stored in a DMAP record, representing either a scalar or /// vector field. -#[derive(Debug, Clone, PartialEq, FromPyObject, IntoPyObject)] +#[derive(Debug, Clone, PartialEq)] #[repr(C)] pub enum DmapField { Vector(DmapVec), @@ -449,14 +381,6 @@ impl DmapField { } } } -// impl IntoPyObject for DmapField { -// fn into_py(self, py: Python<'_>) -> PyObject { -// match self { -// DmapField::Scalar(x) => x.into_py(py), -// DmapField::Vector(x) => x.into_py(py), -// } -// } -// } /// Macro for implementing conversion traits between primitives and [`DmapField`], [`DmapScalar`] /// types. @@ -957,7 +881,7 @@ pub fn check_vector_opt( #[cfg(test)] mod tests { use super::*; - use numpy::ndarray::array; + use ndarray::array; #[test] fn dmaptype() -> Result<()> { @@ -1191,7 +1115,7 @@ mod tests { #[test] fn check_fields_in_indexmap() -> Result<()> { - use numpy::ndarray::array; + use ndarray::array; let mut rec = IndexMap::::new(); let res = check_scalar(&rec, "test", &Type::Char);