Skip to main content

datui_cli/
units.rs

1//! Sizes and durations written with their unit: `512MiB`, `250ms`. One parser reads
2//! them from a config file, from `-c` and from a flag, so a value means the same
3//! wherever it is written.
4
5use std::time::Duration;
6
7/// Binary units, and the decimal ones people also write. A size has a unit unless it
8/// is 0: a bare `512` could be bytes or the MiB an older key took.
9const SIZE_UNITS: &[(&str, u64)] = &[
10    ("B", 1),
11    ("KiB", 1 << 10),
12    ("MiB", 1 << 20),
13    ("GiB", 1 << 30),
14    ("TiB", 1 << 40),
15    ("kB", 1_000),
16    ("KB", 1_000),
17    ("MB", 1_000_000),
18    ("GB", 1_000_000_000),
19    ("TB", 1_000_000_000_000),
20];
21
22const DURATION_UNITS: &[(&str, u64)] = &[("ms", 1), ("s", 1_000), ("m", 60_000), ("h", 3_600_000)];
23
24/// The number and the unit of `text`, the unit's factor looked up in `units`.
25fn split(
26    text: &str,
27    units: &[(&str, u64)],
28    what: &str,
29    example: &str,
30) -> Result<(f64, u64), String> {
31    let text = text.trim();
32    let at = text
33        .find(|c: char| !(c.is_ascii_digit() || c == '.' || c == '_'))
34        .unwrap_or(text.len());
35    let (number, unit) = text.split_at(at);
36    let number = number.replace('_', "");
37    let value: f64 = number
38        .parse()
39        .map_err(|_| format!("\"{text}\" is not a {what}, such as {example}"))?;
40    let unit = unit.trim();
41    if unit.is_empty() {
42        if value == 0.0 {
43            return Ok((0.0, 1));
44        }
45        return Err(format!("\"{text}\" needs a unit, as in {example}"));
46    }
47    let factor = units
48        .iter()
49        .find(|(name, _)| *name == unit)
50        .map(|(_, f)| *f)
51        .ok_or_else(|| {
52            let names: Vec<&str> = units.iter().map(|(n, _)| *n).collect();
53            format!("\"{unit}\" is not a unit of {what}: {}", names.join(", "))
54        })?;
55    Ok((value, factor))
56}
57
58/// Bytes from `512MiB`, `2GiB`, `1.5GB` or `0`.
59pub fn parse_size(text: &str) -> Result<u64, String> {
60    let (value, factor) = split(text, SIZE_UNITS, "size", "512MiB")?;
61    let bytes = value * factor as f64;
62    if !bytes.is_finite() || bytes > u64::MAX as f64 {
63        return Err(format!("\"{}\" is too large", text.trim()));
64    }
65    Ok(bytes.round() as u64)
66}
67
68/// A duration from `250ms`, `1.5s`, `2m` or `0`.
69pub fn parse_duration(text: &str) -> Result<Duration, String> {
70    let (value, factor) = split(text, DURATION_UNITS, "duration", "250ms")?;
71    let ms = value * factor as f64;
72    if !ms.is_finite() || ms > u64::MAX as f64 {
73        return Err(format!("\"{}\" is too long", text.trim()));
74    }
75    Ok(Duration::from_millis(ms.round() as u64))
76}
77
78/// `bytes` in the largest binary unit that holds it whole: `512MiB`, `100KiB`.
79pub fn format_size(bytes: u64) -> String {
80    if bytes == 0 {
81        return "0".to_string();
82    }
83    let (name, factor) = [
84        ("TiB", 1u64 << 40),
85        ("GiB", 1 << 30),
86        ("MiB", 1 << 20),
87        ("KiB", 1 << 10),
88    ]
89    .into_iter()
90    .find(|(_, f)| bytes.is_multiple_of(*f))
91    .unwrap_or(("B", 1));
92    format!("{}{name}", bytes / factor)
93}
94
95/// `duration` in the largest unit that holds it whole: `250ms`, `2s`, `1m`.
96pub fn format_duration(duration: Duration) -> String {
97    let ms = u64::try_from(duration.as_millis()).unwrap_or(u64::MAX);
98    if ms == 0 {
99        return "0".to_string();
100    }
101    let (name, factor) = DURATION_UNITS
102        .iter()
103        .rev()
104        .find(|(_, f)| ms.is_multiple_of(*f))
105        .copied()
106        .unwrap_or(("ms", 1));
107    format!("{}{name}", ms / factor)
108}
109
110/// A size in bytes, written with its unit in config: `max_buffered = "512MiB"`.
111#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)]
112pub struct ByteSize(pub u64);
113
114impl ByteSize {
115    pub const fn mib(n: u64) -> Self {
116        Self(n << 20)
117    }
118    pub const fn kib(n: u64) -> Self {
119        Self(n << 10)
120    }
121    pub fn bytes(self) -> u64 {
122        self.0
123    }
124}
125
126impl std::fmt::Display for ByteSize {
127    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
128        f.write_str(&format_size(self.0))
129    }
130}
131
132impl std::str::FromStr for ByteSize {
133    type Err = String;
134    fn from_str(s: &str) -> Result<Self, String> {
135        parse_size(s).map(Self)
136    }
137}
138
139/// A duration written with its unit in config: `follow_interval = "250ms"`.
140#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)]
141pub struct Interval(pub Duration);
142
143impl Interval {
144    pub const fn ms(n: u64) -> Self {
145        Self(Duration::from_millis(n))
146    }
147    pub fn duration(self) -> Duration {
148        self.0
149    }
150}
151
152impl std::fmt::Display for Interval {
153    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
154        f.write_str(&format_duration(self.0))
155    }
156}
157
158impl std::str::FromStr for Interval {
159    type Err = String;
160    fn from_str(s: &str) -> Result<Self, String> {
161        parse_duration(s).map(Self)
162    }
163}
164
165macro_rules! serde_as_text {
166    ($ty:ty) => {
167        impl serde::Serialize for $ty {
168            fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
169                s.serialize_str(&self.to_string())
170            }
171        }
172
173        impl<'de> serde::Deserialize<'de> for $ty {
174            fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
175                struct Visit;
176                impl serde::de::Visitor<'_> for Visit {
177                    type Value = $ty;
178                    fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
179                        f.write_str("a value with its unit, in quotes")
180                    }
181                    fn visit_str<E: serde::de::Error>(self, v: &str) -> Result<$ty, E> {
182                        v.parse().map_err(E::custom)
183                    }
184                    // `0` needs no unit or quotes.
185                    fn visit_i64<E: serde::de::Error>(self, v: i64) -> Result<$ty, E> {
186                        self.visit_str(&v.to_string())
187                    }
188                    fn visit_u64<E: serde::de::Error>(self, v: u64) -> Result<$ty, E> {
189                        self.visit_str(&v.to_string())
190                    }
191                }
192                d.deserialize_any(Visit)
193            }
194        }
195    };
196}
197
198serde_as_text!(ByteSize);
199serde_as_text!(Interval);
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204
205    #[test]
206    fn sizes_read_and_write_with_their_unit() {
207        assert_eq!(parse_size("512MiB"), Ok(512 << 20));
208        assert_eq!(parse_size(" 2GiB "), Ok(2 << 30));
209        assert_eq!(parse_size("1.5KiB"), Ok(1536));
210        assert_eq!(parse_size("100KB"), Ok(100_000));
211        assert_eq!(parse_size("0"), Ok(0));
212        assert!(parse_size("512").unwrap_err().contains("needs a unit"));
213        assert!(parse_size("5mb").unwrap_err().contains("not a unit"));
214        assert!(parse_size("MiB").is_err());
215        assert_eq!(format_size(512 << 20), "512MiB");
216        assert_eq!(format_size(2048 << 20), "2GiB");
217        assert_eq!(format_size(100 << 10), "100KiB");
218        assert_eq!(format_size(1000), "1000B");
219        assert_eq!(format_size(0), "0");
220    }
221
222    #[test]
223    fn durations_read_and_write_with_their_unit() {
224        assert_eq!(parse_duration("250ms"), Ok(Duration::from_millis(250)));
225        assert_eq!(parse_duration("1.5s"), Ok(Duration::from_millis(1500)));
226        assert_eq!(parse_duration("2m"), Ok(Duration::from_secs(120)));
227        assert_eq!(parse_duration("0"), Ok(Duration::ZERO));
228        assert!(parse_duration("250").unwrap_err().contains("needs a unit"));
229        assert_eq!(format_duration(Duration::from_millis(1500)), "1500ms");
230        assert_eq!(format_duration(Duration::from_secs(2)), "2s");
231        assert_eq!(format_duration(Duration::from_secs(60)), "1m");
232    }
233}