1use regex::Regex;
2use serde::{Deserialize, Serialize};
3
4#[derive(Debug, Clone, Default)]
5pub struct ByteSizeConf(pub u64);
6
7impl ByteSizeConf {
8 pub fn as_bytes(&self) -> u64 {
9 self.0
10 }
11
12 pub fn as_kilobytes(&self) -> u64 {
13 self.0 / 1024
14 }
15
16 pub fn as_megabytes(&self) -> u64 {
17 self.0 / 1024 / 1024
18 }
19
20 pub fn as_gigabytes(&self) -> u64 {
21 self.0 / 1024 / 1024 / 1024
22 }
23
24 fn to_string(&self) -> String {
25 let size = self.0;
26 if size < 1024 {
27 format!("{} bytes", size)
28 } else if size < 2u64.pow(20) {
29 format!("{} KiB", size / 1024)
30 } else if size < 2u64.pow(30) {
31 format!("{} MiB", size / 2u64.pow(20))
32 } else {
33 format!("{} GiB", size / 2u64.pow(30))
34 }
35 }
36}
37
38impl<'de> Deserialize<'de> for ByteSizeConf {
39 fn deserialize<D>(deserializer: D) -> Result<ByteSizeConf, D::Error>
40 where
41 D: serde::Deserializer<'de>,
42 {
43 let s = String::deserialize(deserializer)?;
44 let size = parse_byte_size(&s).map_err(serde::de::Error::custom)?;
45 Ok(ByteSizeConf(size))
46 }
47}
48
49impl Serialize for ByteSizeConf {
50 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
51 where
52 S: serde::Serializer,
53 {
54 self.to_string().serialize(serializer)
55 }
56}
57
58impl std::fmt::Display for ByteSizeConf {
59 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
60 write!(f, "{}", self.to_string())
61 }
62}
63
64fn parse_byte_size(s: &str) -> Result<u64, String> {
65 let re = Regex::new(r"^(\d+)\s*(b|kb|ki|mb|mi|gb|gi|tb|ti)$").map_err(|_| "Invalid regex")?;
66 let s = s.to_lowercase();
67 let caps = re
68 .captures(s.as_str())
69 .ok_or_else(|| format!("Invalid byte size: {}", s))?;
70 let size = caps
71 .get(1)
72 .ok_or_else(|| "Invalid byte size".to_string())?
73 .as_str()
74 .parse()
75 .map_err(|_| "Invalid byte size".to_string())?;
76 let unit = caps
77 .get(2)
78 .ok_or_else(|| "Invalid byte size".to_string())?
79 .as_str()
80 .to_lowercase();
81 let unit = unit.as_str();
82
83 match unit {
84 "b" => Ok(size),
85 "kb" => Ok(size * 1000),
86 "ki" => Ok(size * 2u64.pow(10)),
87 "mb" => Ok(size * 10u64.pow(6)),
88 "mi" => Ok(size * 2u64.pow(20)),
89 "gb" => Ok(size * 10u64.pow(9)),
90 "gi" => Ok(size * 2u64.pow(30)),
91 "tb" => Ok(size * 10u64.pow(12)),
92 "ti" => Ok(size * 2u64.pow(40)),
93 _ => Err("Invalid byte size".to_string()),
94 }
95}
96
97#[cfg(test)]
98mod tests {
99 use super::*;
100
101 #[test]
102 fn test_parse_byte_size() {
103 assert_eq!(parse_byte_size("1b").unwrap(), 1);
104 assert_eq!(parse_byte_size("1kb").unwrap(), 1000);
105 assert_eq!(parse_byte_size("1ki").unwrap(), 1024);
106 assert_eq!(parse_byte_size("1mb").unwrap(), 1000 * 1000);
107 assert_eq!(parse_byte_size("1mi").unwrap(), 1024 * 1024);
108 assert_eq!(parse_byte_size("1gb").unwrap(), 1000 * 1000 * 1000);
109 assert_eq!(parse_byte_size("1gi").unwrap(), 1024 * 1024 * 1024);
110 assert_eq!(parse_byte_size("1tb").unwrap(), 1000 * 1000 * 1000 * 1000);
111 assert_eq!(parse_byte_size("1ti").unwrap(), 1024 * 1024 * 1024 * 1024);
112 }
113
114 #[test]
115 fn test_byte_size_display() {
116 assert_eq!(ByteSizeConf(1).to_string(), "1 bytes");
117 assert_eq!(ByteSizeConf(1024).to_string(), "1 KiB");
118 assert_eq!(ByteSizeConf(1024 * 1024).to_string(), "1 MiB");
119 assert_eq!(ByteSizeConf(1024 * 1024 * 1024).to_string(), "1 GiB");
120 }
121}