Merge pull request #27 from SuperDARNCanada/develop

Release: v0.4.0

Author: RemingtonRohel

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "dmap"
-version = "0.3.0"
+version = "0.4.0"
 edition = "2021"
 rust-version = "1.63.0"
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "darn-dmap"
-version = "0.3.0"
+version = "0.4.0"
 requires-python = ">=3.8"
 authors = [
     { name = "Remington Rohel" }

src/compression.rs

@@ -0,0 +1,45 @@
+use std::io::{Chain, Cursor, Error, Read};
+
+/// Detects bz2 compression on the input `stream`. Returns a reader
+/// which includes all data from `stream`.
+pub(crate) fn detect_bz2<T>(mut stream: T) -> Result<(bool, Chain<Cursor<[u8; 3]>, T>), Error>
+where
+    T: for<'a> Read,
+{
+    // Read the first 3 bytes to detect bz2 compression
+    let mut buffer = [0u8; 3];
+    stream.read_exact(&mut buffer)?;
+
+    // valid bz2 blocks start with "BZh", which is 425a68 in hex.
+    let is_bz2 = buffer == [0x42, 0x5a, 0x68];
+    let full_stream = Cursor::new(buffer).chain(stream);
+    Ok((is_bz2, full_stream))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use bzip2::{read::BzDecoder, read::BzEncoder, Compression};
+
+    #[test]
+    fn bz2_detection() -> Result<(), Error> {
+        let data = "Hello world".as_bytes();
+        let compressor = BzEncoder::new(data, Compression::best());
+
+        let (result, stream) = detect_bz2(compressor)?;
+        assert_eq!(result, true);
+        let mut returned_stream = vec![];
+        let mut decompressed = BzDecoder::new(stream);
+        let _ = decompressed.read_to_end(&mut returned_stream);
+        assert_eq!(returned_stream, b"Hello world");
+
+        let data = "Hello world".as_bytes();
+        let (result, mut stream) = detect_bz2(data)?;
+        assert_eq!(result, false);
+        let mut returned_stream = vec![];
+        let _ = stream.read_to_end(&mut returned_stream);
+        assert_eq!(returned_stream, b"Hello world");
+
+        Ok(())
+    }
+}

src/error.rs

@@ -14,6 +14,10 @@ pub enum DmapError {
     #[error("{0}")]
     Io(#[from] std::io::Error),
 
+    /// Error casting between Dmap types.
+    #[error("{0}")]
+    BadCast(#[from] std::num::TryFromIntError),
+
     /// Invalid key for a DMAP type. Valid keys are defined [here](https://github.com/SuperDARN/rst/blob/main/codebase/general/src.lib/dmap.1.25/include/dmap.h)
     #[error("{0}")]
     InvalidKey(i8),
@@ -31,6 +35,10 @@ pub enum DmapError {
     #[error("{0}")]
     InvalidVector(String),
 
+    /// Bytes cannot be interpreted as a DMAP field.
+    #[error("{0}")]
+    InvalidField(String),
+
     /// Errors when reading in multiple records
     #[error("First error: {1}\nRecords with errors: {0:?}")]
     BadRecords(Vec<usize>, String),

src/lib.rs

@@ -5,20 +5,21 @@
 //! 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/).
 
+pub mod compression;
 pub mod error;
 pub mod formats;
 pub mod record;
 pub mod types;
 
-use crate::error::DmapError;
-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;
+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 bzip2::read::BzEncoder;
 use bzip2::Compression;
@@ -234,7 +235,14 @@ macro_rules! read_py {
     }
 }
 
-read_py!(iqdat, "read_iqdat", "read_iqdat_lax", "read_iqdat_bytes", "read_iqdat_bytes_lax", "sniff_iqdat");
+read_py!(
+    iqdat,
+    "read_iqdat",
+    "read_iqdat_lax",
+    "read_iqdat_bytes",
+    "read_iqdat_bytes_lax",
+    "sniff_iqdat"
+);
 read_py!(
     rawacf,
     "read_rawacf",
@@ -251,10 +259,38 @@ read_py!(
     "read_fitacf_bytes_lax",
     "sniff_fitacf"
 );
-read_py!(grid, "read_grid", "read_grid_lax", "read_grid_bytes", "read_grid_bytes_lax", "sniff_grid");
-read_py!(map, "read_map", "read_map_lax", "read_map_bytes", "read_map_bytes_lax", "sniff_map");
-read_py!(snd, "read_snd", "read_snd_lax", "read_snd_bytes", "read_snd_bytes_lax", "sniff_snd");
-read_py!(dmap, "read_dmap", "read_dmap_lax", "read_dmap_bytes", "read_dmap_bytes_lax", "sniff_dmap");
+read_py!(
+    grid,
+    "read_grid",
+    "read_grid_lax",
+    "read_grid_bytes",
+    "read_grid_bytes_lax",
+    "sniff_grid"
+);
+read_py!(
+    map,
+    "read_map",
+    "read_map_lax",
+    "read_map_bytes",
+    "read_map_bytes_lax",
+    "sniff_map"
+);
+read_py!(
+    snd,
+    "read_snd",
+    "read_snd_lax",
+    "read_snd_bytes",
+    "read_snd_bytes_lax",
+    "sniff_snd"
+);
+read_py!(
+    dmap,
+    "read_dmap",
+    "read_dmap_lax",
+    "read_dmap_bytes",
+    "read_dmap_bytes_lax",
+    "sniff_dmap"
+);
 
 /// Checks that a list of dictionaries contains DMAP records, then appends to outfile.
 ///

src/record.rs

@@ -1,12 +1,12 @@
 //! Defines the `Record` trait, which contains the shared behaviour that all
 //! DMAP records must have.
 
+use crate::compression::detect_bz2;
 use crate::error::DmapError;
 use crate::types::{parse_scalar, parse_vector, read_data, DmapField, DmapType, DmapVec, Fields};
 use bzip2::read::BzDecoder;
 use indexmap::IndexMap;
 use rayon::prelude::*;
-use std::ffi::OsStr;
 use std::fmt::Debug;
 use std::fs::File;
 use std::io::{Cursor, Read};
@@ -26,12 +26,19 @@ pub trait Record<'a>:
         Self: Sized,
         Self: Send,
     {
-        let mut buffer = [0; 8];  // record size should be an i32 of the data
-        let read_result = dmap_data.read(&mut buffer[..])?;
-        if read_result < buffer.len() {
-            return Err(DmapError::CorruptStream("Unable to read size of first record"))
+        let mut stream: Box<dyn Read>;
+        let (is_bz2, chunk) = detect_bz2(&mut dmap_data)?;
+        if is_bz2 {
+            stream = Box::new(BzDecoder::new(chunk));
+        } else {
+            stream = Box::new(chunk);
         }
 
+        let mut buffer = [0; 8]; // record size should be an i32 of the data
+        stream
+            .read_exact(&mut buffer)
+            .map_err(|_| DmapError::CorruptStream("Unable to read size of first record"))?;
+
         let rec_size = i32::from_le_bytes(buffer[4..8].try_into().unwrap()) as usize; // advance 4 bytes, skipping the "code" field
         if rec_size <= 0 {
             return Err(DmapError::InvalidRecord(format!(
@@ -42,7 +49,7 @@ pub trait Record<'a>:
 
         let mut rec = vec![0; rec_size];
         rec[0..8].clone_from_slice(&buffer[..]);
-        dmap_data.read_exact(&mut rec[8..])?;
+        stream.read_exact(&mut rec[8..])?;
         let first_rec = Self::parse_record(&mut Cursor::new(rec))?;
 
         Ok(first_rec)
@@ -57,12 +64,19 @@ pub trait Record<'a>:
         Self: Send,
     {
         let mut buffer: Vec<u8> = vec![];
-        dmap_data.read_to_end(&mut buffer)?;
+        let (is_bz2, mut chunk) = detect_bz2(&mut dmap_data)?;
+        if is_bz2 {
+            let mut stream = BzDecoder::new(chunk);
+            stream.read_to_end(&mut buffer)?;
+        } else {
+            chunk.read_to_end(&mut buffer)?;
+        }
 
         let mut slices: Vec<_> = vec![];
         let mut rec_start: usize = 0;
         let mut rec_size: usize;
         let mut rec_end: usize;
+
         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
@@ -123,7 +137,13 @@ pub trait Record<'a>:
         Self: Send,
     {
         let mut buffer: Vec<u8> = vec![];
-        dmap_data.read_to_end(&mut buffer)?;
+        let (is_bz2, mut chunk) = detect_bz2(&mut dmap_data)?;
+        if is_bz2 {
+            let mut stream = BzDecoder::new(chunk);
+            stream.read_to_end(&mut buffer)?;
+        } else {
+            chunk.read_to_end(&mut buffer)?;
+        }
 
         let mut dmap_records: Vec<Self> = vec![];
         let mut bad_byte: Option<usize> = None;
@@ -173,13 +193,7 @@ pub trait Record<'a>:
         Self: Send,
     {
         let file = File::open(infile)?;
-        match infile.extension() {
-            Some(ext) if ext == OsStr::new("bz2") => {
-                let compressor = BzDecoder::new(file);
-                Self::read_records(compressor)
-            }
-            _ => Self::read_records(file),
-        }
+        Self::read_records(file)
     }
 
     /// Read a DMAP file of type `Self`.
@@ -192,13 +206,7 @@ pub trait Record<'a>:
         Self: Send,
     {
         let file = File::open(infile)?;
-        match infile.extension() {
-            Some(ext) if ext == OsStr::new("bz2") => {
-                let compressor = BzDecoder::new(file);
-                Self::read_records_lax(compressor)
-            }
-            _ => Self::read_records_lax(file),
-        }
+        Self::read_records_lax(file)
     }
 
     /// Reads the first record of a DMAP file of type `Self`.
@@ -208,13 +216,7 @@ pub trait Record<'a>:
         Self: Send,
     {
         let file = File::open(infile)?;
-        match infile.extension() {
-            Some(ext) if ext == OsStr::new("bz2") => {
-                let compressor = BzDecoder::new(file);
-                Self::read_first_record(compressor)
-            }
-            _ => Self::read_first_record(file),
-        }
+        Self::read_first_record(file)
     }
 
     /// Reads a record from `cursor`.
@@ -648,6 +650,45 @@ macro_rules! create_record_type {
                     Self::coerce::<[< $format:camel Record>]>(value, &$fields)
                 }
             }
+
+            #[cfg(test)]
+            mod tests {
+                use super::*;
+                use std::path::PathBuf;
+
+                /// Creates a test to ensure that the record is still able to be read, even when missing
+                /// some of the optional fields.
+                #[test]
+                fn test_missing_optional_fields() -> Result<(), DmapError> {
+                    let filename: PathBuf = PathBuf::from(format!("tests/test_files/test.{}", stringify!($format)));
+                    let data = [< $format:camel Record >]::sniff_file(&filename).expect("Unable to sniff file");
+                    let recs = data.inner();
+
+                    for field in $fields.scalars_optional.iter().chain($fields.vectors_optional.iter()) {
+                        let mut cloned_rec = recs.clone();
+                        let _ = cloned_rec.shift_remove(field.0);
+                        let _ = [< $format:camel Record >]::try_from(&mut cloned_rec)?;
+                    }
+                    Ok(())
+                }
+
+                /// Creates a test to ensure that the record is not able to be read when missing
+                /// some of the required fields.
+                #[test]
+                fn test_missing_required_fields() -> Result<(), DmapError> {
+                    let filename: PathBuf = PathBuf::from(format!("tests/test_files/test.{}", stringify!($format)));
+                    let data = [< $format:camel Record >]::sniff_file(&filename).expect("Unable to sniff file");
+                    let recs = data.inner();
+
+                    for field in $fields.scalars_required.iter().chain($fields.vectors_required.iter()) {
+                        let mut cloned_rec = recs.clone();
+                        let _ = cloned_rec.shift_remove(field.0);
+                        let res = [< $format:camel Record >]::try_from(&mut cloned_rec);
+                        assert!(res.is_err());
+                    }
+                    Ok(())
+                }
+            }
         }
     }
 }