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

relay_config/
byte_size.rs
