relay_config/
byte_size.rs

1use std::fmt;
2use std::str::FromStr;
3
4pub use human_size::ParsingError as ByteSizeParseError;
5use human_size::{Any, Size, SpecificSize};
6use serde::ser::Serializer;
7use serde::{Serialize, de};
8
9/// Represents a size in bytes.
10///
11/// `ByteSize` can be parsed from strings or with Serde, and remembers the original unit that was
12/// used to describe it for stable serialization. Use `ByteSize::infer` to infer the most
13/// appropriate unit for a number of bytes.
14///
15/// Units based on 1000 and 1024 are both supported:
16///  - To refer to the 1000-based versions, use "kB" and "MB".
17///  - To refer to the 1024-based versions, use "KiB" and "MiB".
18///
19/// # Examples
20///
21/// Infer the best unit:
22///
23/// ```
24/// use relay_config::ByteSize;
25///
26/// let size = ByteSize::infer(42 * 1000 * 1000);
27/// assert_eq!("42MB", size.to_string());
28/// ```
29///
30/// Format a 1024-based size to string:
31///
32/// ```
33/// use relay_config::ByteSize;
34///
35/// let size = ByteSize::kibibytes(42);
36/// assert_eq!("42KiB", size.to_string());
37/// ```
38pub struct ByteSize(Size);
39
40impl ByteSize {
41    fn multiple(value: u32, multiple: Any) -> Self {
42        // Can be unwrapped because f64::from<u32> always returns a "normal" number.
43        // See https://doc.rust-lang.org/nightly/std/primitive.f64.html#method.is_normal
44        Self(SpecificSize::new(value, multiple).unwrap())
45    }
46
47    fn try_multiple(value: u32, multiple: Any) -> Option<Self> {
48        let factor = match multiple {
49            Any::Mebibyte => 1024 * 1024,
50            Any::Megabyte => 1000 * 1000,
51            Any::Kibibyte => 1024,
52            Any::Kilobyte => 1000,
53            _ => 1,
54        };
55
56        match value % factor {
57            0 => Some(Self::multiple(value / factor, multiple)),
58            _ => None,
59        }
60    }
61
62    /// Create a byte size from bytes, inferring the most appropriate unit.
63    pub fn infer(value: u32) -> Self {
64        Self::try_multiple(value, Any::Mebibyte)
65            .or_else(|| Self::try_multiple(value, Any::Megabyte))
66            .or_else(|| Self::try_multiple(value, Any::Kibibyte))
67            .or_else(|| Self::try_multiple(value, Any::Kilobyte))
68            .unwrap_or_else(|| Self::bytes(value))
69    }
70
71    /// Create a byte size from bytes.
72    pub fn bytes(value: u32) -> Self {
73        Self::multiple(value, Any::Byte)
74    }
75
76    /// Create a byte size from 1024-based kibibytes.
77    pub fn kibibytes(value: u32) -> Self {
78        Self::multiple(value, Any::Kibibyte)
79    }
80
81    /// Create a byte size from 1024-based mebibytes.
82    pub fn mebibytes(value: u32) -> Self {
83        Self::multiple(value, Any::Mebibyte)
84    }
85
86    /// Return the value in bytes.
87    pub fn as_bytes(&self) -> usize {
88        let byte_size = self.0.into::<human_size::Byte>();
89        byte_size.value() as usize
90    }
91}
92
93impl From<u32> for ByteSize {
94    fn from(value: u32) -> ByteSize {
95        ByteSize::infer(value)
96    }
97}
98
99impl FromStr for ByteSize {
100    type Err = ByteSizeParseError;
101
102    fn from_str(value: &str) -> Result<Self, Self::Err> {
103        match value.parse::<u32>() {
104            Ok(bytes) => Ok(Self::bytes(bytes)),
105            Err(_) => value.parse().map(ByteSize),
106        }
107    }
108}
109
110impl fmt::Display for ByteSize {
111    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
112        write!(f, "{}{}", self.0.value(), self.0.multiple())
113    }
114}
115
116impl fmt::Debug for ByteSize {
117    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
118        f.debug_tuple("ByteSize")
119            .field(&format_args!("{}", self.0))
120            .finish()
121    }
122}
123
124impl Serialize for ByteSize {
125    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
126    where
127        S: Serializer,
128    {
129        serializer.collect_str(self)
130    }
131}
132
133impl<'de> de::Deserialize<'de> for ByteSize {
134    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
135    where
136        D: de::Deserializer<'de>,
137    {
138        struct V;
139
140        impl de::Visitor<'_> for V {
141            type Value = ByteSize;
142
143            fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
144                formatter.write_str("data size")
145            }
146
147            fn visit_u64<E>(self, value: u64) -> Result<ByteSize, E>
148            where
149                E: de::Error,
150            {
151                match value.try_into() {
152                    Ok(value32) => Ok(ByteSize::infer(value32)),
153                    Err(_) => Err(de::Error::invalid_value(
154                        de::Unexpected::Unsigned(value),
155                        &self,
156                    )),
157                }
158            }
159
160            fn visit_str<E>(self, value: &str) -> Result<ByteSize, E>
161            where
162                E: de::Error,
163            {
164                value
165                    .parse()
166                    .map_err(|_| de::Error::invalid_value(de::Unexpected::Str(value), &self))
167            }
168        }
169
170        deserializer.deserialize_any(V)
171    }
172}
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177
178    #[test]
179    fn test_infer() {
180        let size = ByteSize::infer(42);
181        assert_eq!(42, size.as_bytes());
182        assert_eq!("42B", size.to_string());
183
184        let size = ByteSize::infer(1000);
185        assert_eq!(1000, size.as_bytes());
186        assert_eq!("1kB", size.to_string());
187
188        let size = ByteSize::infer(1024);
189        assert_eq!(1024, size.as_bytes());
190        assert_eq!("1KiB", size.to_string());
191
192        let size = ByteSize::infer(1000 * 1000);
193        assert_eq!(1000 * 1000, size.as_bytes());
194        assert_eq!("1MB", size.to_string());
195
196        let size = ByteSize::infer(1024 * 1024);
197        assert_eq!(1024 * 1024, size.as_bytes());
198        assert_eq!("1MiB", size.to_string());
199    }
200
201    #[test]
202    fn test_parse() {
203        let size = ByteSize::from_str("4242").unwrap();
204        assert_eq!(4242, size.as_bytes());
205        assert_eq!("4242B", size.to_string());
206
207        let size = ByteSize::from_str("42B").unwrap();
208        assert_eq!(42, size.as_bytes());
209        assert_eq!("42B", size.to_string());
210
211        // NOTE: Lowercase k is kilo
212        let size = ByteSize::from_str("1kB").unwrap();
213        assert_eq!(1000, size.as_bytes());
214        assert_eq!("1kB", size.to_string());
215
216        // NOTE: Uppercase K is kibi
217        let size = ByteSize::from_str("1KB").unwrap();
218        assert_eq!(1024, size.as_bytes());
219        assert_eq!("1KiB", size.to_string());
220
221        let size = ByteSize::from_str("1KiB").unwrap();
222        assert_eq!(1024, size.as_bytes());
223        assert_eq!("1KiB", size.to_string());
224
225        let size = ByteSize::from_str("1MB").unwrap();
226        assert_eq!(1000 * 1000, size.as_bytes());
227        assert_eq!("1MB", size.to_string());
228
229        let size = ByteSize::from_str("1MiB").unwrap();
230        assert_eq!(1024 * 1024, size.as_bytes());
231        assert_eq!("1MiB", size.to_string());
232    }
233
234    #[test]
235    fn test_as_bytes() {
236        let size = ByteSize::bytes(42);
237        assert_eq!(42, size.as_bytes());
238
239        let size = ByteSize::kibibytes(42);
240        assert_eq!(42 * 1024, size.as_bytes());
241
242        let size = ByteSize::mebibytes(42);
243        assert_eq!(42 * 1024 * 1024, size.as_bytes());
244    }
245
246    #[test]
247    fn test_serde_number() {
248        let size = serde_json::from_str::<ByteSize>("1024").unwrap();
249        let json = serde_json::to_string(&size).unwrap();
250        assert_eq!(json, "\"1KiB\"");
251    }
252
253    #[test]
254    fn test_serde_string() {
255        let size = serde_json::from_str::<ByteSize>("\"1KiB\"").unwrap();
256        let json = serde_json::to_string(&size).unwrap();
257        assert_eq!(json, "\"1KiB\"");
258    }
259}
relay_config/byte_size.rs

relay_config/
byte_size.rs
