tocat_api/
size.rs

1//! size.rs: one grammar for every byte count tocat accepts.
2//!
3//! Shared rather than per-crate because `buffer-size` in the config, `size=` on
4//! a pipe endpoint and `bytes=` on the `limit` plugin should not disagree about
5//! what `10M` means. Suffixes are binary: `k` is 1024, not 1000, because the
6//! things being sized are buffers and transfers.
7
8use std::{fmt, str::FromStr};
9
10use serde::{Deserialize, Serialize};
11
12/// A byte count.
13///
14/// Accepts a plain number or a suffix: `65536`, `64k`, `1MiB`.
15#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
16pub struct ByteSize(pub usize);
17
18impl ByteSize {
19    #[must_use]
20    pub fn bytes(self) -> usize {
21        self.0
22    }
23}
24
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct ParseSizeError(String);
27
28impl fmt::Display for ParseSizeError {
29    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
30        f.write_str(&self.0)
31    }
32}
33
34impl std::error::Error for ParseSizeError {}
35
36impl fmt::Display for ByteSize {
37    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
38        const UNITS: [(usize, &str); 3] = [
39            (1024 * 1024 * 1024, "GiB"),
40            (1024 * 1024, "MiB"),
41            (1024, "KiB"),
42        ];
43
44        for (scale, suffix) in UNITS {
45            if self.0 >= scale && self.0.is_multiple_of(scale) {
46                return write!(f, "{}{suffix}", self.0 / scale);
47            }
48        }
49
50        write!(f, "{}", self.0)
51    }
52}
53
54impl FromStr for ByteSize {
55    type Err = ParseSizeError;
56
57    fn from_str(raw: &str) -> Result<Self, Self::Err> {
58        let trimmed = raw.trim();
59        let digits = trimmed
60            .trim_end_matches(|c: char| c.is_ascii_alphabetic())
61            .trim_end();
62        let suffix = trimmed[digits.len()..].trim().to_ascii_lowercase();
63
64        let scale: usize = match suffix.as_str() {
65            "" | "b" => 1,
66            "k" | "kb" | "kib" => 1024,
67            "m" | "mb" | "mib" => 1024 * 1024,
68            "g" | "gb" | "gib" => 1024 * 1024 * 1024,
69            other => {
70                return Err(ParseSizeError(format!(
71                    "unknown size suffix {other:?}; use k, m or g"
72                )));
73            }
74        };
75
76        let value: usize = digits
77            .parse()
78            .map_err(|_| ParseSizeError(format!("{digits:?} is not a number")))?;
79
80        let bytes = value
81            .checked_mul(scale)
82            .ok_or_else(|| ParseSizeError(format!("{raw} overflows a size")))?;
83
84        Ok(ByteSize(bytes))
85    }
86}
87
88impl Serialize for ByteSize {
89    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
90        serializer.serialize_str(&self.to_string())
91    }
92}
93
94impl<'de> Deserialize<'de> for ByteSize {
95    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
96        use serde::de::Error as _;
97
98        // Both `buffer-size = 262144` and `buffer-size = "256KiB"` are natural
99        // things to write, so accept either.
100        #[derive(Deserialize)]
101        #[serde(untagged)]
102        enum Raw {
103            Bytes(usize),
104            Text(String),
105        }
106
107        match Raw::deserialize(deserializer)? {
108            Raw::Bytes(n) => Ok(ByteSize(n)),
109            Raw::Text(s) => s.parse().map_err(D::Error::custom),
110        }
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    #[test]
119    fn suffixes_are_binary() {
120        assert_eq!("64k".parse::<ByteSize>().unwrap(), ByteSize(65536));
121        assert_eq!("1MiB".parse::<ByteSize>().unwrap(), ByteSize(1024 * 1024));
122        assert_eq!("512".parse::<ByteSize>().unwrap(), ByteSize(512));
123    }
124
125    #[test]
126    fn round_multiples_display_with_a_suffix() {
127        assert_eq!(ByteSize(1024 * 1024).to_string(), "1MiB");
128        assert_eq!(ByteSize(1500).to_string(), "1500");
129    }
130
131    #[test]
132    fn nonsense_is_rejected() {
133        assert!("10furlongs".parse::<ByteSize>().is_err());
134        assert!("k".parse::<ByteSize>().is_err());
135    }
136}
tocat_api/size.rs

tocat_api/
size.rs
