Skip to main content

ecoord_io/
util.rs

1use std::io::{Read, Write};
2use strum_macros::EnumIter;
3
4pub const DEFAULT_COMPRESSION_LEVEL: i32 = 10;
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, EnumIter)]
7pub enum Compression {
8    #[default]
9    None,
10    Zstd(i32),
11}
12
13impl Compression {
14    pub fn as_str(&self) -> Option<&'static str> {
15        match self {
16            Compression::None => None,
17            Compression::Zstd(_) => Some("zst"),
18        }
19    }
20
21    pub fn from_str(s: &str) -> Option<Self> {
22        match s {
23            "zst" => Some(Compression::Zstd(DEFAULT_COMPRESSION_LEVEL)),
24            _ => None,
25        }
26    }
27}
28
29impl Compression {
30    pub const fn default_zstd() -> Self {
31        Self::Zstd(DEFAULT_COMPRESSION_LEVEL)
32    }
33
34    pub fn is_enabled(&self) -> bool {
35        !matches!(self, Compression::None)
36    }
37
38    pub fn level(&self) -> Option<i32> {
39        match self {
40            Compression::None => None,
41            Compression::Zstd(level) => Some(*level),
42        }
43    }
44}
45
46impl Compression {
47    pub fn wrap_reader<'a, R: Read + 'a>(
48        &self,
49        reader: R,
50    ) -> Result<Box<dyn Read + 'a>, std::io::Error> {
51        match self {
52            Compression::None => Ok(Box::new(reader)),
53            Compression::Zstd(_) => {
54                let decoder = zstd::Decoder::new(reader)?;
55                Ok(Box::new(decoder))
56            }
57        }
58    }
59
60    pub fn wrap_writer<'a, W: Write + 'a>(
61        &self,
62        writer: W,
63    ) -> Result<Box<dyn Write + 'a>, std::io::Error> {
64        match self {
65            Compression::None => Ok(Box::new(writer)),
66            Compression::Zstd(level) => {
67                let encoder = zstd::Encoder::new(writer, *level)?;
68                Ok(Box::new(encoder.auto_finish()))
69            }
70        }
71    }
72}