pub fn parse_size(input: &str) -> Result<u64, String> {
let s = input.trim();
if s.is_empty() {
return Err(format!("empty size literal: {input:?}"));
}
let split_at = s
.find(|c: char| !c.is_ascii_digit())
.unwrap_or(s.len());
let (num_part, unit_part) = s.split_at(split_at);
if num_part.is_empty() {
return Err(format!("size literal {input:?} has no number"));
}
let n: u64 = num_part
.parse()
.map_err(|_| format!("size literal {input:?} has invalid number: {num_part:?}"))?;
let multiplier = parse_unit(unit_part.trim())
.ok_or_else(|| format!("size literal {input:?} has unknown unit: {unit_part:?}"))?;
n.checked_mul(multiplier)
.ok_or_else(|| format!("size literal {input:?} overflows u64"))
}
fn parse_unit(s: &str) -> Option<u64> {
match s.to_ascii_lowercase().as_str() {
"" | "b" => Some(1),
"k" | "kb" => Some(1024),
"m" | "mb" => Some(1024 * 1024),
"g" | "gb" => Some(1024 * 1024 * 1024),
"t" | "tb" => Some(1024 * 1024 * 1024 * 1024),
_ => None,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn bare_integer_is_bytes() {
assert_eq!(parse_size("0").unwrap(), 0);
assert_eq!(parse_size("512").unwrap(), 512);
assert_eq!(parse_size("1024").unwrap(), 1024);
}
#[test]
fn binary_multipliers_match_redis() {
assert_eq!(parse_size("1k").unwrap(), 1024);
assert_eq!(parse_size("1kb").unwrap(), 1024);
assert_eq!(parse_size("1KB").unwrap(), 1024);
assert_eq!(parse_size("64mb").unwrap(), 64 * 1024 * 1024);
assert_eq!(parse_size("2gb").unwrap(), 2 * 1024 * 1024 * 1024);
assert_eq!(parse_size("1tb").unwrap(), 1024_u64.pow(4));
}
#[test]
fn whitespace_tolerated() {
assert_eq!(parse_size(" 64 mb ").unwrap(), 64 * 1024 * 1024);
assert_eq!(parse_size("64 mb").unwrap(), 64 * 1024 * 1024);
}
#[test]
fn empty_or_no_number_rejected() {
assert!(parse_size("").is_err());
assert!(parse_size(" ").is_err());
assert!(parse_size("mb").is_err());
}
#[test]
fn unknown_unit_rejected() {
assert!(parse_size("64xb").is_err());
assert!(parse_size("64 zz").is_err());
}
#[test]
fn float_rejected() {
assert!(parse_size("1.5gb").is_err());
}
#[test]
fn overflow_reported() {
let huge = format!("{}kb", u64::MAX / 1024 + 1);
assert!(parse_size(&huge).is_err());
}
}