Skip to main content

bytesize/
parse.rs

1use alloc::{borrow::ToOwned as _, format, string::String};
2use core::{fmt, str};
3
4use super::ByteSize;
5
6impl str::FromStr for ByteSize {
7    type Err = String;
8
9    fn from_str(value: &str) -> Result<Self, Self::Err> {
10        if let Ok(v) = value.parse::<u64>() {
11            return Ok(Self(v));
12        }
13        let number = take_while(value, |c| c.is_ascii_digit() || c == '.');
14        match number.parse::<f64>() {
15            Ok(v) => {
16                let suffix = skip_while(&value[number.len()..], char::is_whitespace);
17                match suffix.parse::<Unit>() {
18                    // Use exact integer arithmetic when the number has no
19                    // fractional part. `f64` only has a 53-bit mantissa, so byte
20                    // counts at or above 2^53 would otherwise be rounded (e.g.
21                    // "9007199254740993B" parsed to 9007199254740992).
22                    Ok(u) if !number.contains('.') => match number.parse::<u64>() {
23                        Ok(n) => Ok(Self(n.saturating_mul(u.factor()))),
24                        Err(_) => Ok(Self((v * u) as u64)),
25                    },
26                    Ok(u) => Ok(Self((v * u) as u64)),
27                    Err(error) => Err(format!(
28                        "couldn't parse {suffix:?} into a known SI unit, {error}"
29                    )),
30                }
31            }
32            Err(error) => Err(format!("couldn't parse {value:?} into a ByteSize, {error}")),
33        }
34    }
35}
36
37fn take_while<P>(s: &str, mut predicate: P) -> &str
38where
39    P: FnMut(char) -> bool,
40{
41    let offset = s
42        .chars()
43        .take_while(|ch| predicate(*ch))
44        .map(|ch| ch.len_utf8())
45        .sum();
46    &s[..offset]
47}
48
49fn skip_while<P>(s: &str, mut predicate: P) -> &str
50where
51    P: FnMut(char) -> bool,
52{
53    let offset: usize = s
54        .chars()
55        .skip_while(|ch| predicate(*ch))
56        .map(|ch| ch.len_utf8())
57        .sum();
58    &s[(s.len() - offset)..]
59}
60
61/// Scale unit.
62///
63/// ```
64/// use bytesize::Unit;
65///
66/// assert_eq!(
67///     "GiB".parse::<Unit>().unwrap(),
68///     Unit::GibiByte,
69/// );
70///
71/// "gibibyte".parse::<Unit>().unwrap_err();
72/// ```
73#[non_exhaustive]
74#[derive(Debug, Clone, PartialEq)]
75pub enum Unit {
76    /// Single byte.
77    Byte,
78
79    // power of tens
80    /// Kilobyte (10^3 bytes).
81    KiloByte,
82
83    /// Megabyte (10^6 bytes)
84    MegaByte,
85
86    /// Gigabyte (10^9 bytes)
87    GigaByte,
88
89    /// Terabyte (10^12 bytes)
90    TeraByte,
91
92    /// Petabyte (10^15 bytes)
93    PetaByte,
94
95    /// Exabyte (10^18 bytes)
96    ExaByte,
97
98    // power of twos
99    /// Kibibyte (2^10 bytes)
100    KibiByte,
101
102    /// Mebibyte (2^20 bytes)
103    MebiByte,
104
105    /// Gibibyte (2^30 bytes)
106    GibiByte,
107
108    /// Tebibyte (2^40 bytes)
109    TebiByte,
110
111    /// Pebibyte (2^50 bytes)
112    PebiByte,
113
114    /// Exbibyte (2^60 bytes)
115    ExbiByte,
116}
117
118impl Unit {
119    fn factor(&self) -> u64 {
120        match self {
121            Self::Byte => 1,
122            // decimal units
123            Self::KiloByte => crate::KB,
124            Self::MegaByte => crate::MB,
125            Self::GigaByte => crate::GB,
126            Self::TeraByte => crate::TB,
127            Self::PetaByte => crate::PB,
128            Self::ExaByte => crate::EB,
129            // binary units
130            Self::KibiByte => crate::KIB,
131            Self::MebiByte => crate::MIB,
132            Self::GibiByte => crate::GIB,
133            Self::TebiByte => crate::TIB,
134            Self::PebiByte => crate::PIB,
135            Self::ExbiByte => crate::EIB,
136        }
137    }
138}
139
140mod impl_ops {
141    use super::Unit;
142    use core::ops;
143
144    impl ops::Add<u64> for Unit {
145        type Output = u64;
146
147        fn add(self, other: u64) -> Self::Output {
148            self.factor() + other
149        }
150    }
151
152    impl ops::Add<Unit> for u64 {
153        type Output = u64;
154
155        fn add(self, other: Unit) -> Self::Output {
156            self + other.factor()
157        }
158    }
159
160    impl ops::Mul<u64> for Unit {
161        type Output = u64;
162
163        fn mul(self, other: u64) -> Self::Output {
164            self.factor() * other
165        }
166    }
167
168    impl ops::Mul<Unit> for u64 {
169        type Output = u64;
170
171        fn mul(self, other: Unit) -> Self::Output {
172            self * other.factor()
173        }
174    }
175
176    impl ops::Add<f64> for Unit {
177        type Output = f64;
178
179        fn add(self, other: f64) -> Self::Output {
180            self.factor() as f64 + other
181        }
182    }
183
184    impl ops::Add<Unit> for f64 {
185        type Output = f64;
186
187        fn add(self, other: Unit) -> Self::Output {
188            other.factor() as f64 + self
189        }
190    }
191
192    impl ops::Mul<f64> for Unit {
193        type Output = f64;
194
195        fn mul(self, other: f64) -> Self::Output {
196            self.factor() as f64 * other
197        }
198    }
199
200    impl ops::Mul<Unit> for f64 {
201        type Output = f64;
202
203        fn mul(self, other: Unit) -> Self::Output {
204            other.factor() as f64 * self
205        }
206    }
207}
208
209impl str::FromStr for Unit {
210    type Err = UnitParseError;
211
212    fn from_str(unit: &str) -> Result<Self, Self::Err> {
213        match () {
214            _ if unit.eq_ignore_ascii_case("b") => Ok(Self::Byte),
215            _ if unit.eq_ignore_ascii_case("k") | unit.eq_ignore_ascii_case("kb") => {
216                Ok(Self::KiloByte)
217            }
218            _ if unit.eq_ignore_ascii_case("m") | unit.eq_ignore_ascii_case("mb") => {
219                Ok(Self::MegaByte)
220            }
221            _ if unit.eq_ignore_ascii_case("g") | unit.eq_ignore_ascii_case("gb") => {
222                Ok(Self::GigaByte)
223            }
224            _ if unit.eq_ignore_ascii_case("t") | unit.eq_ignore_ascii_case("tb") => {
225                Ok(Self::TeraByte)
226            }
227            _ if unit.eq_ignore_ascii_case("p") | unit.eq_ignore_ascii_case("pb") => {
228                Ok(Self::PetaByte)
229            }
230            _ if unit.eq_ignore_ascii_case("e") | unit.eq_ignore_ascii_case("eb") => {
231                Ok(Self::ExaByte)
232            }
233            _ if unit.eq_ignore_ascii_case("ki") | unit.eq_ignore_ascii_case("kib") => {
234                Ok(Self::KibiByte)
235            }
236            _ if unit.eq_ignore_ascii_case("mi") | unit.eq_ignore_ascii_case("mib") => {
237                Ok(Self::MebiByte)
238            }
239            _ if unit.eq_ignore_ascii_case("gi") | unit.eq_ignore_ascii_case("gib") => {
240                Ok(Self::GibiByte)
241            }
242            _ if unit.eq_ignore_ascii_case("ti") | unit.eq_ignore_ascii_case("tib") => {
243                Ok(Self::TebiByte)
244            }
245            _ if unit.eq_ignore_ascii_case("pi") | unit.eq_ignore_ascii_case("pib") => {
246                Ok(Self::PebiByte)
247            }
248            _ if unit.eq_ignore_ascii_case("ei") | unit.eq_ignore_ascii_case("eib") => {
249                Ok(Self::ExbiByte)
250            }
251            _ => Err(UnitParseError(to_string_truncate(unit))),
252        }
253    }
254}
255
256/// Safely truncates
257fn to_string_truncate(unit: &str) -> String {
258    const MAX_UNIT_LEN: usize = 3;
259
260    if unit.len() > MAX_UNIT_LEN {
261        // TODO(MSRV 1.91): use ceil_char_boundary
262
263        if unit.is_char_boundary(3) {
264            format!("{}...", &unit[..3])
265        } else if unit.is_char_boundary(4) {
266            format!("{}...", &unit[..4])
267        } else if unit.is_char_boundary(5) {
268            format!("{}...", &unit[..5])
269        } else if unit.is_char_boundary(6) {
270            format!("{}...", &unit[..6])
271        } else {
272            unreachable!("char boundary will be within 4 bytes")
273        }
274    } else {
275        unit.to_owned()
276    }
277}
278
279/// Error returned when parsing a [`Unit`] fails.
280#[derive(Debug)]
281pub struct UnitParseError(String);
282
283impl fmt::Display for UnitParseError {
284    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
285        write!(f, "Failed to parse unit \"{}\"", self.0)
286    }
287}
288
289#[cfg(feature = "std")]
290impl std::error::Error for UnitParseError {}
291
292#[cfg(test)]
293mod tests {
294    use alloc::string::ToString as _;
295
296    use super::*;
297
298    #[test]
299    fn truncating_error_strings() {
300        assert_eq!("", to_string_truncate(""));
301        assert_eq!("b", to_string_truncate("b"));
302        assert_eq!("ob", to_string_truncate("ob"));
303        assert_eq!("foo", to_string_truncate("foo"));
304
305        assert_eq!("foo...", to_string_truncate("foob"));
306        assert_eq!("foo...", to_string_truncate("foobar"));
307    }
308
309    #[test]
310    fn when_ok() {
311        // shortcut for writing test cases
312        fn parse(s: &str) -> u64 {
313            s.parse::<ByteSize>().unwrap().0
314        }
315
316        assert_eq!("0".parse::<ByteSize>().unwrap().0, 0);
317        assert_eq!(parse("0"), 0);
318        assert_eq!(parse("500"), 500);
319        assert_eq!(parse("1K"), Unit::KiloByte * 1);
320        assert_eq!(parse("1Ki"), Unit::KibiByte * 1);
321        assert_eq!(parse("1.5Ki"), (1.5 * Unit::KibiByte) as u64);
322        assert_eq!(parse("1KiB"), 1 * Unit::KibiByte);
323        assert_eq!(parse("1.5KiB"), (1.5 * Unit::KibiByte) as u64);
324        assert_eq!(parse("3 MB"), Unit::MegaByte * 3);
325        assert_eq!(parse("4 MiB"), Unit::MebiByte * 4);
326        assert_eq!(parse("6 GB"), 6 * Unit::GigaByte);
327        assert_eq!(parse("4 GiB"), 4 * Unit::GibiByte);
328        assert_eq!(parse("88TB"), 88 * Unit::TeraByte);
329        assert_eq!(parse("521TiB"), 521 * Unit::TebiByte);
330        assert_eq!(parse("8 PB"), 8 * Unit::PetaByte);
331        assert_eq!(parse("8P"), 8 * Unit::PetaByte);
332        assert_eq!(parse("12 PiB"), 12 * Unit::PebiByte);
333    }
334
335    #[test]
336    fn large_integer_byte_counts_are_exact() {
337        // shortcut for writing test cases
338        fn parse(s: &str) -> u64 {
339            s.parse::<ByteSize>().unwrap().0
340        }
341
342        // 2^53 + 1 is not exactly representable as an `f64`, so the unit-suffix
343        // parsing path must not round it down. Regression test for the
344        // `f64`-based parsing path.
345        assert_eq!(parse("9007199254740993B"), 9_007_199_254_740_993);
346        assert_eq!(parse("9007199254740993"), parse("9007199254740993B"));
347    }
348
349    #[test]
350    fn when_err() {
351        // shortcut for writing test cases
352        fn parse(s: &str) -> Result<ByteSize, String> {
353            s.parse::<ByteSize>()
354        }
355
356        assert!(parse("").is_err());
357        assert!(parse("a124GB").is_err());
358        assert!(parse("1.3 42.0 B").is_err());
359        assert!(parse("1.3 ... B").is_err());
360        // The original implementation did not account for the possibility that users may
361        // use whitespace to visually separate digits, thus treat it as an error
362        assert!(parse("1 000 B").is_err());
363    }
364
365    #[test]
366    fn to_and_from_str() {
367        // shortcut for writing test cases
368        fn parse(s: &str) -> u64 {
369            s.parse::<ByteSize>().unwrap().0
370        }
371
372        assert_eq!(parse(&parse("128GB").to_string()), 128 * Unit::GigaByte);
373        assert_eq!(
374            parse(&ByteSize(parse("128.000 GiB")).to_string()),
375            128 * Unit::GibiByte,
376        );
377    }
378}