deserialize-file-size/src/lib.rs

//! A serde helper function for deserializing file size input
//! flexibly and robustly.
//!
//! Accepts either:
//!
//! 1) a "human" size string, e.g. "1k", "5mb", "12GiB", etc.
//! 2) an integer number of bytes
//!
//! # Examples
//!
//! ```
//! use serde::Deserialize;
//! use deserialize_file_size::deserialize_file_size;
//!
//! #[derive(Deserialize, Debug, PartialEq)]
//! struct FileSize {
//!     #[serde(deserialize_with = "deserialize_file_size")]
//!     sz: usize,
//! }
//!
//! let size_str = r#"{"sz": "42mb"}"#;
//! assert_eq!(
//!     serde_json::from_str::<FileSize>(size_str).unwrap(),
//!     FileSize { sz: 1024 * 1024 * 42 },
//! );
//!
//! let int_bytes = r#"{"sz": 4096}"#;
//! assert_eq!(
//!     serde_json::from_str::<FileSize>(int_bytes).unwrap(),
//!     FileSize { sz: 4096 },
//! );
//! ```

#![allow(clippy::single_char_pattern)] // annoying, who the f cares "a" vs 'a'

use std::fmt;
use serde::de::{self, Visitor};

/// returns size in bytes if parsing is successful. note: units "k"/"kb", "m"/"mb",
/// and "g"/"gb" are converted to KiB, MiB, and GiB respectively. there are no 1000-based
/// file sizes as far as this implementation is concerned, 1024 is file size god.
pub fn parse_file_size(s: &str) -> Option<usize> {
    let mut s = s.trim().to_string();
    s[..].make_ascii_lowercase();
    let s = if s.contains("ib") {
        s
    } else {
        s.replace("kb", "k")        // first truncate kb/mb/gb -> k/m/g
            .replace("mb", "m")
            .replace("gb", "g")
            .replace("k", "kib")    // then transform k/m/g -> kib/mib/gib
            .replace("m", "mib")
            .replace("g", "gib")
    };

    let bytes: u128 = byte_unit::Byte::from_str(&s)
        .ok()?
        .get_bytes();

    usize::try_from(bytes).ok()
}

/// A serde "deserialize_with" helper function for parsing a `usize` field 
/// flexibly and robustly.
///
/// Accepts input that is either:
///
/// 1) a "human" size string, e.g. "10k", "42mb", "7GiB", etc.
/// 2) an integer number of bytes
///
/// To make the point explicit, either `String`/`&str` or integer
/// (`visit_u64`-compatible) input is accepted.
///
/// # Examples
///
/// ```
/// use serde::Deserialize;
/// use deserialize_file_size::deserialize_file_size;
///
/// #[derive(Deserialize, Debug, PartialEq)]
/// struct FileSize {
///     #[serde(deserialize_with = "deserialize_file_size")]
///     sz: usize,
/// }
///
/// let size_str = r#"{"sz": "42mb"}"#;
/// assert_eq!(
///     serde_json::from_str::<FileSize>(size_str).unwrap(),
///     FileSize { sz: 1024 * 1024 * 42 },
/// );
///
/// let int_bytes = r#"{"sz": 4096}"#;
/// assert_eq!(
///     serde_json::from_str::<FileSize>(int_bytes).unwrap(),
///     FileSize { sz: 4096 },
/// );
/// ```
pub fn deserialize_file_size<'de, D>(deserializer: D) -> Result<usize, D::Error>
    where D: serde::de::Deserializer<'de>
{
    struct SizeStringOrUsize;

    impl<'de> Visitor<'de> for SizeStringOrUsize {
        type Value = usize;

        fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
            formatter.write_str("file size string or integer number of bytes")
        }

        fn visit_str<E>(self, value: &str) -> Result<Self::Value, E>
        where
            E: de::Error,
        {
            match parse_file_size(value) {
                Some(size) => Ok(size),
                None => {
                    Err(serde::de::Error::custom(format!(
                        "parsing file size input '{}' failed; expected number followed \
                         by human size abbreviation (e.g. 10k/5mb/15GiB) or integer number \
                         of bytes",
                         value,
                    )))
                }
            }
        }

        fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E>
            where E: de::Error,
        {
            Ok(value as usize)
        }
    }

    deserializer.deserialize_any(SizeStringOrUsize)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn check_file_size_parser_against_several_inputs() {
        assert_eq!(parse_file_size("1k")    , Some(1024));
        assert_eq!(parse_file_size("1kib")  , Some(1024));
        assert_eq!(parse_file_size("1K")    , Some(1024));
        assert_eq!(parse_file_size("1Kb")   , Some(1024));
        assert_eq!(parse_file_size("1kb")   , Some(1024));
        assert_eq!(parse_file_size("1 kb")  , Some(1024));
        assert_eq!(parse_file_size("1 k")   , Some(1024));
        assert_eq!(parse_file_size("1 kib") , Some(1024));

        assert_eq!(parse_file_size("1m")    , Some(1024 * 1024));
        assert_eq!(parse_file_size("1mib")  , Some(1024 * 1024));
        assert_eq!(parse_file_size("1M")    , Some(1024 * 1024));
        assert_eq!(parse_file_size("1Mb")   , Some(1024 * 1024));
        assert_eq!(parse_file_size("1MB")   , Some(1024 * 1024));
        assert_eq!(parse_file_size("1mb")   , Some(1024 * 1024));
        assert_eq!(parse_file_size("1 MB")  , Some(1024 * 1024));

        assert_eq!(parse_file_size("1g")    , Some(1024 * 1024 * 1024));
        assert_eq!(parse_file_size("1gib")  , Some(1024 * 1024 * 1024));
        assert_eq!(parse_file_size("1G")    , Some(1024 * 1024 * 1024));
        assert_eq!(parse_file_size("1Gb")   , Some(1024 * 1024 * 1024));
        assert_eq!(parse_file_size("1gb")   , Some(1024 * 1024 * 1024));
        assert_eq!(parse_file_size("1 gb")  , Some(1024 * 1024 * 1024));
        assert_eq!(parse_file_size("1 g")   , Some(1024 * 1024 * 1024));
        assert_eq!(parse_file_size("1 Gib") , Some(1024 * 1024 * 1024));

        assert_eq!(parse_file_size("48G") , Some( 48 * 1024 * 1024 * 1024));
        assert_eq!(parse_file_size("96G") , Some( 96 * 1024 * 1024 * 1024));
        assert_eq!(parse_file_size("2G")  , Some(  2 * 1024 * 1024 * 1024));
        assert_eq!(parse_file_size("128G"), Some(128 * 1024 * 1024 * 1024));
    }

    #[test]
    fn check_deserialize_file_size() {
        #[derive(serde::Deserialize)]
        struct A {
            #[serde(deserialize_with = "deserialize_file_size")]
            sz: usize,
        }

        assert_eq!(
            serde_json::from_str::<A>(r#"{"sz":"1mb"}"#).unwrap().sz,
            1024 * 1024,
        );
        assert_eq!(
            serde_json::from_str::<A>(r#"{"sz":4096}"#).unwrap().sz,
            4096,
        );
    }
}