dmap/

lib.rs

//! A library for SuperDARN DMAP file I/O.
//!
//! [![github]](https://github.com/SuperDARNCanada/dmap) [![crates-io]](https://crates.io/crates/darn-dmap) [![docs-rs]](crate)
//!
//! [github]: https://img.shields.io/badge/github-8da0cb?style=for-the-badge&labelColor=555555&logo=github
//! [crates-io]: https://img.shields.io/badge/crates.io-fc8d62?style=for-the-badge&labelColor=555555&logo=rust
//! [docs-rs]: https://img.shields.io/badge/docs.rs-66c2a5?style=for-the-badge&labelColor=555555&logo=docs.rs
//!
//! <br>
//!
//! This library also has a Python API using pyo3.
//!
//! For more information about DMAP files, see [RST](https://radar-software-toolkit-rst.readthedocs.io/en/latest/)
//! or [pyDARNio](https://pydarnio.readthedocs.io/en/latest/).
//!
//! The main feature of this crate is the [`Record`] trait, which defines a valid DMAP record and functions for
//! reading from and writing to byte streams. The SuperDARN file formats IQDAT, RAWACF, FITACF, GRID, MAP, and SND are
//! all supported with structs that implement [`Record`], namely:
//!
//! - [`IqdatRecord`]
//! - [`RawacfRecord`]
//! - [`FitacfRecord`]
//! - [`GridRecord`]
//! - [`MapRecord`]
//! - [`SndRecord`]
//!
//! Each struct has a list of required and optional fields that it uses to verify the integrity of the record.
//! Only fields listed in the required and optional lists are allowed, no required field can be missing, and
//! each field has an expected primitive type. Additionally, each format has groupings of vector fields which
//! must all share the same shape; e.g. `acfd` and `xcfd` in a RAWACF file.
//!
//! There is also a generic [`DmapRecord`] struct which has no knowledge of required or optional fields. When reading from
//! a byte stream, the parsed data will be identical when using both a specific format like [`RawacfRecord`] and the generic
//! [`DmapRecord`]; however, when writing to a byte stream, the output may differ. Since [`DmapRecord`] has no knowledge of
//! the expected primitive type for each field, it defaults to a type that fits the data. For example, the `stid` field may
//! be saved as an `i8` when using [`DmapRecord`] instead of an `i16` which [`RawacfRecord`] specifies it must be.
//!
//! <div class="note">
//! Each type of record has a specific field ordering hard-coded by this library. This is the order in which fields are written to file,
//! and may not match the ordering of fields generated by RST. This also means that round-trip I/O (i.e. reading a file and
//! writing back out to a new file) is not guaranteed to generate an identical file; however, it is guaranteed that all the
//! information is the same, just not necessarily in the same order.
//! </div>
//!
//! <br>
//!
//! # Examples
//!
//! Convenience functions for reading from and writing to a file exist to simplify the most common use cases.
//! This is defined by [`Record::read_file`]
//! ```
//! use dmap::*;
//! use std::path::PathBuf;
//!
//! # fn main() -> Result<(), DmapError> {
//! let path = PathBuf::from("tests/test_files/test.rawacf");
//! let rawacf_data = RawacfRecord::read_file(&path)?;
//! let unchecked_data = DmapRecord::read_file(&path)?;
//!
//! assert_eq!(rawacf_data.len(), unchecked_data.len());
//! assert_eq!(rawacf_data[0].get(&"stid".to_string()), unchecked_data[0].get(&"stid".to_string()));
//!
//! // Write the records to a file
//! let out_path = PathBuf::from("tests/test_files/output.rawacf");
//! RawacfRecord::write_to_file(&rawacf_data, &out_path, false)?;
//! # std::fs::remove_file(out_path)?;
//! #    Ok(())
//! # }
//! ```
//! You can read from anything that implements the `Read` trait using the functions exposed by the [`Record`] trait.
//! Detection and decompression of bz2 is also conducted automatically.
//! ```
//! use dmap::*;
//! use std::fs::File;
//! use itertools::izip;
//!
//! # fn main() -> Result<(), DmapError> {
//! let file = File::open("tests/test_files/test.rawacf.bz2")?;  // `File` implements the `Read` trait
//! let rawacf_data = RawacfRecord::read_records(file)?;
//!
//! let uncompressed_data = RawacfRecord::read_file("tests/test_files/test.rawacf")?;
//! assert_eq!(rawacf_data.len(), uncompressed_data.len());
//! for (left, right) in izip!(rawacf_data, uncompressed_data) {
//!     assert_eq!(left, right)
//! }
//! #     Ok(())
//! # }
//! ```

pub(crate) mod compression;
pub mod error;
pub mod formats;
pub(crate) mod io;
pub mod record;
pub mod types;

pub use crate::error::DmapError;
pub use crate::formats::dmap::DmapRecord;
pub use crate::formats::fitacf::FitacfRecord;
pub use crate::formats::grid::GridRecord;
pub use crate::formats::iqdat::IqdatRecord;
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<IndexMap>` 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 >]<P: AsRef<Path>>(
                recs: Vec<IndexMap<String, DmapField>>,
                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`
/// * `sniff_[name]` - reads only the first record from file.
/// * `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,
        $metadata_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<Vec<IndexMap<String, DmapField>>> {
                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<IndexMap<String, DmapField>>, Option<usize>)> {
                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<Vec<IndexMap<String, DmapField>>> {
                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<IndexMap<String, DmapField>>, Option<usize>)> {
                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 first record." ]
            #[pyfunction]
            #[pyo3(name = $sniff_name)]
            #[pyo3(text_signature = "(infile: str, /)")]
            fn [< sniff_ $name _py >](infile: PathBuf) -> PyResult<IndexMap<String, DmapField>> {
                Ok([< $name:camel Record >]::sniff_file(&infile)
                    .map_err(PyErr::from)?
                    .inner()
                )
            }

            #[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<Vec<IndexMap<String, DmapField>>> {
                Ok([< $name:camel Record >]::read_file_metadata(&infile)
                    .map_err(PyErr::from)?
                )
            }
        }
    }
}

read_py!(
    iqdat,
    "read_iqdat",
    "read_iqdat_lax",
    "read_iqdat_bytes",
    "read_iqdat_bytes_lax",
    "sniff_iqdat",
    "read_iqdat_metadata"
);
read_py!(
    rawacf,
    "read_rawacf",
    "read_rawacf_lax",
    "read_rawacf_bytes",
    "read_rawacf_bytes_lax",
    "sniff_rawacf",
    "read_rawacf_metadata"
);
read_py!(
    fitacf,
    "read_fitacf",
    "read_fitacf_lax",
    "read_fitacf_bytes",
    "read_fitacf_bytes_lax",
    "sniff_fitacf",
    "read_fitacf_metadata"
);
read_py!(
    grid,
    "read_grid",
    "read_grid_lax",
    "read_grid_bytes",
    "read_grid_bytes_lax",
    "sniff_grid",
    "read_grid_metadata"
);
read_py!(
    map,
    "read_map",
    "read_map_lax",
    "read_map_bytes",
    "read_map_bytes_lax",
    "sniff_map",
    "read_map_metadata"
);
read_py!(
    snd,
    "read_snd",
    "read_snd_lax",
    "read_snd_bytes",
    "read_snd_bytes_lax",
    "sniff_snd",
    "read_snd_metadata"
);
read_py!(
    dmap,
    "read_dmap",
    "read_dmap_lax",
    "read_dmap_bytes",
    "read_dmap_bytes_lax",
    "sniff_dmap",
    "read_dmap_metadata"
);

/// 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<IndexMap<String, DmapField>>,
    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<IndexMap<String, DmapField>>,
    bz2: bool,
) -> PyResult<Py<PyAny>> {
    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<IndexMap<String, DmapField>>, 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<IndexMap<String, DmapField>>, bz2: bool) -> PyResult<Py<PyAny>> {
                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)?)?;

    // 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 the first record
    m.add_function(wrap_pyfunction!(sniff_dmap_py, m)?)?;
    m.add_function(wrap_pyfunction!(sniff_iqdat_py, m)?)?;
    m.add_function(wrap_pyfunction!(sniff_rawacf_py, m)?)?;
    m.add_function(wrap_pyfunction!(sniff_fitacf_py, m)?)?;
    m.add_function(wrap_pyfunction!(sniff_snd_py, m)?)?;
    m.add_function(wrap_pyfunction!(sniff_grid_py, m)?)?;
    m.add_function(wrap_pyfunction!(sniff_map_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)?)?;

    Ok(())
}