k8s_quantity_parser/
lib.rs

1#![warn(clippy::all, clippy::pedantic, clippy::nursery, clippy::cargo)]
2use eyre::{eyre, Report};
3use k8s_openapi::apimachinery::pkg::api::resource::Quantity;
4use regex::Regex;
5
6#[allow(non_camel_case_types)]
7enum QuantityMemoryUnits {
8    Ki,
9    Mi,
10    Gi,
11    Ti,
12    Pi,
13    Ei,
14    k,
15    M,
16    G,
17    T,
18    P,
19    E,
20    m,
21    Invalid,
22}
23
24impl QuantityMemoryUnits {
25    fn new(unit: &str) -> Self {
26        match unit {
27            "Ki" => Self::Ki,
28            "Mi" => Self::Mi,
29            "Gi" => Self::Gi,
30            "Ti" => Self::Ti,
31            "Pi" => Self::Pi,
32            "Ei" => Self::Ei,
33            "k" => Self::k,
34            "M" => Self::M,
35            "G" => Self::G,
36            "T" => Self::T,
37            "P" => Self::P,
38            "E" => Self::E,
39            "m" => Self::m,
40            _ => Self::Invalid,
41        }
42    }
43}
44
45/// This trait works as a parser for the values retrieved from BTreeMap<String, Quantity> collections
46/// in `k8s_openapi::api::core::v1::Pod` and `k8s_openapi::api::core::v1::Node`
47///
48/// # Errors
49/// The parser will fails if encounters an invalid unit letters or failed to parse String to i64
50
51pub trait QuantityParser {
52    /// This method will parse the cpu resource values returned by Kubernetes Api
53    ///
54    /// ```rust
55    /// # use k8s_openapi::apimachinery::pkg::api::resource::Quantity;
56    /// # use k8s_quantity_parser::QuantityParser;
57    /// #
58    /// let mib = Quantity("1Mi".into());
59    /// let ret: i64 = 1048576;
60    /// assert_eq!(mib.to_bytes().ok().flatten().unwrap(), ret);
61    /// ```
62    ///
63    /// # Errors
64    ///
65    /// The parser will fails if encounters an invalid unit letters or failed to parse String to i64
66    ///
67    fn to_milli_cpus(&self) -> Result<Option<i64>, Report>;
68    /// This method will parse the memory resource values returned by Kubernetes Api
69    ///
70    /// ```rust
71    /// # use k8s_openapi::apimachinery::pkg::api::resource::Quantity;
72    /// # use k8s_quantity_parser::QuantityParser;
73    /// #
74    /// let cpu = Quantity("4".into());
75    /// let ret: i64 = 4000;
76    /// assert_eq!(cpu.to_milli_cpus().ok().flatten().unwrap(), ret)
77    /// ```
78    ///
79    /// # Errors
80    ///
81    /// The parser will fails if encounters an invalid unit letters or failed to parse String to i64
82    ///
83    fn to_bytes(&self) -> Result<Option<i64>, Report>;
84}
85
86impl QuantityParser for Quantity {
87    fn to_milli_cpus(&self) -> Result<Option<i64>, Report> {
88        let unit_str = &self.0;
89        let rgx = Regex::new(r"([m]{1}$)")?;
90        let cap = rgx.captures(unit_str);
91        if cap.is_none() {
92            return Ok(Some(unit_str.parse::<i64>()? * 1000));
93        };
94        let mt = cap.unwrap().get(0).unwrap();
95        let unit_str = unit_str.replace(mt.as_str(), "");
96        Ok(Some(unit_str.parse::<i64>()?))
97    }
98
99    fn to_bytes(&self) -> Result<Option<i64>, Report> {
100        let unit_str = &self.0;
101        let rgx = Regex::new(r"([[:alpha:]]{1,2}$)")?;
102        let cap = rgx.captures(unit_str);
103
104        if cap.is_none() {
105            return Ok(Some(unit_str.parse::<i64>()?));
106        };
107
108        // Is safe to use unwrap here, as the value is already checked.
109        match cap.unwrap().get(0) {
110            Some(m) => match QuantityMemoryUnits::new(m.as_str()) {
111                QuantityMemoryUnits::Ki => {
112                    let unit_str = unit_str.replace(m.as_str(), "");
113                    let amount = unit_str.parse::<i64>()?;
114                    Ok(Some(amount * 1024))
115                }
116                QuantityMemoryUnits::Mi => {
117                    let unit_str = unit_str.replace(m.as_str(), "");
118                    let amount = unit_str.parse::<i64>()?;
119                    Ok(Some((amount * 1024) * 1024))
120                }
121                QuantityMemoryUnits::Gi => {
122                    let unit_str = unit_str.replace(m.as_str(), "");
123                    let amount = unit_str.parse::<i64>()?;
124                    Ok(Some(((amount * 1024) * 1024) * 1024))
125                }
126                QuantityMemoryUnits::Ti => {
127                    let unit_str = unit_str.replace(m.as_str(), "");
128                    let amount = unit_str.parse::<i64>()?;
129                    Ok(Some((((amount * 1024) * 1024) * 1024) * 1024))
130                }
131                QuantityMemoryUnits::Pi => {
132                    let unit_str = unit_str.replace(m.as_str(), "");
133                    let amount = unit_str.parse::<i64>()?;
134                    Ok(Some(((((amount * 1024) * 1024) * 1024) * 1024) * 1024))
135                }
136                QuantityMemoryUnits::Ei => {
137                    let unit_str = unit_str.replace(m.as_str(), "");
138                    let amount = unit_str.parse::<i64>()?;
139                    Ok(Some(
140                        (((((amount * 1024) * 1024) * 1024) * 1024) * 1024) * 1024,
141                    ))
142                }
143                QuantityMemoryUnits::k => {
144                    let unit_str = unit_str.replace(m.as_str(), "");
145                    let amount = unit_str.parse::<i64>()?;
146                    Ok(Some(amount * 1000))
147                }
148                QuantityMemoryUnits::M => {
149                    let unit_str = unit_str.replace(m.as_str(), "");
150                    let amount = unit_str.parse::<i64>()?;
151                    Ok(Some((amount * 1000) * 1000))
152                }
153                QuantityMemoryUnits::G => {
154                    let unit_str = unit_str.replace(m.as_str(), "");
155                    let amount = unit_str.parse::<i64>()?;
156                    Ok(Some(((amount * 1000) * 1000) * 1000))
157                }
158                QuantityMemoryUnits::T => {
159                    let unit_str = unit_str.replace(m.as_str(), "");
160                    let amount = unit_str.parse::<i64>()?;
161                    Ok(Some((((amount * 1000) * 1000) * 1000) * 1000))
162                }
163                QuantityMemoryUnits::P => {
164                    let unit_str = unit_str.replace(m.as_str(), "");
165                    let amount = unit_str.parse::<i64>()?;
166                    Ok(Some(((((amount * 1000) * 1000) * 1000) * 1000) * 1000))
167                }
168                QuantityMemoryUnits::E => {
169                    let unit_str = unit_str.replace(m.as_str(), "");
170                    let amount = unit_str.parse::<i64>()?;
171                    Ok(Some(
172                        (((((amount * 1000) * 1000) * 1000) * 1000) * 1000) * 1000,
173                    ))
174                }
175                QuantityMemoryUnits::m => {
176                    let unit_str = unit_str.replace(m.as_str(), "");
177                    let amount = unit_str.parse::<i64>()?;
178                    Ok(Some(amount / 1000))
179                }
180                QuantityMemoryUnits::Invalid => Err(eyre!("Invalid unit")),
181            },
182            None => Ok(None),
183        }
184    }
185}
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190
191    #[test]
192    fn to_bytes_works() {
193        assert!(Quantity("12345".into()).to_bytes().is_ok())
194    }
195
196    #[test]
197    fn to_bytes_is_some() {
198        assert!(Quantity("12345".into()).to_bytes().unwrap().is_some())
199    }
200
201    #[test]
202    fn to_milli_cpus_works() {
203        assert!(Quantity("12345m".into()).to_milli_cpus().is_ok())
204    }
205
206    #[test]
207    fn to_milli_cpus_is_some() {
208        assert!(Quantity("12345m".into()).to_milli_cpus().unwrap().is_some())
209    }
210
211    #[test]
212    fn invalid_unit_fails() {
213        assert!(Quantity("12345r".into()).to_bytes().is_err())
214    }
215
216    #[test]
217    fn parse_i64_fails() {
218        assert!(Quantity("123.123".into()).to_bytes().is_err())
219    }
220
221    #[test]
222    fn is_none_value() {
223        assert!(Quantity("0Mi".into()).to_bytes().unwrap().is_some())
224    }
225
226    #[test]
227    fn pow2_mb_to_bytes() {
228        let mib = Quantity("1Mi".into());
229        let ret: i64 = 1048576;
230        assert_eq!(mib.to_bytes().ok().flatten().unwrap(), ret);
231    }
232
233    #[test]
234    fn pow10_gb_to_bytes() {
235        let mib = Quantity("1G".into());
236        let ret: i64 = 1000000000;
237        assert_eq!(mib.to_bytes().ok().flatten().unwrap(), ret);
238    }
239
240    #[test]
241    fn cpu_units_value_to_millis() {
242        let cpu = Quantity("1536m".into());
243        let ret: i64 = 1536;
244        assert_eq!(cpu.to_milli_cpus().ok().flatten().unwrap(), ret)
245    }
246
247    #[test]
248    fn cpu_cores_value_to_millis() {
249        let cpu = Quantity("4".into());
250        let ret: i64 = 4000;
251        assert_eq!(cpu.to_milli_cpus().ok().flatten().unwrap(), ret)
252    }
253}