arkham/src/context.rs

use crate::{Command, Opt};
use std::{collections::BTreeMap, env};

#[derive(Clone, Debug, Default)]
pub struct Map(BTreeMap<String, ContextValue>);

impl Map {
    pub fn new() -> Self {
        Map(BTreeMap::new())
    }
    pub fn set(&mut self, name: String, v: ContextValue) {
        self.0.insert(name, v);
    }
    pub fn get(&self, name: &str) -> Option<&ContextValue> {
        self.0.get(name)
    }
}

#[derive(Clone, Debug)]
pub struct Array(Vec<ContextValue>);

#[derive(Clone, Debug)]
pub enum ContextValue {
    String(String),
    Integer(i64),
    Float(f64),
    Bool(bool),
    Map(Map),
    Array(Array),
}

#[derive(Clone, Copy, Debug)]
pub enum ValueKind {
    String,
    Integer,
    Float,
    Bool,
    Map,
    Array,
}

impl ContextValue {
    pub fn string(&self) -> Option<String> {
        if let Self::String(v) = self {
            Some(v.clone())
        } else {
            None
        }
    }

    pub fn bool(&self) -> Option<bool> {
        if let Self::Bool(v) = self {
            Some(*v)
        } else {
            None
        }
    }

    pub fn integer(&self) -> Option<i64> {
        if let Self::Integer(v) = self {
            Some(*v)
        } else {
            None
        }
    }

    pub fn float(&self) -> Option<f64> {
        if let Self::Float(v) = self {
            Some(*v)
        } else {
            None
        }
    }

    pub fn map(&self) -> Map {
        if let Self::Map(v) = self {
            v.clone()
        } else {
            Map::new()
        }
    }

    pub fn array(&self) -> Array {
        if let Self::Array(v) = self {
            v.clone()
        } else {
            Array(vec![])
        }
    }

    pub fn kind(&self) -> ValueKind {
        match self {
            Self::String(_) => ValueKind::String,
            Self::Integer(_) => ValueKind::Integer,
            Self::Float(_) => ValueKind::Float,
            Self::Bool(_) => ValueKind::Bool,
            Self::Map(_) => ValueKind::Map,
            Self::Array(_) => ValueKind::Array,
        }
    }

    pub fn parse_string(value: String, kind: ValueKind) -> Option<Self> {
        match kind {
            ValueKind::String => Some(Self::String(value)),
            ValueKind::Integer => value.parse::<i64>().ok().map(Self::Integer),
            ValueKind::Float => value.parse::<f64>().ok().map(Self::Float),
            ValueKind::Bool => value.parse::<bool>().ok().map(Self::Bool),
            ValueKind::Array => None,
            ValueKind::Map => None,
        }
    }
}

impl From<String> for ContextValue {
    fn from(fr: String) -> Self {
        ContextValue::String(fr)
    }
}

impl From<i64> for ContextValue {
    fn from(fr: i64) -> Self {
        ContextValue::Integer(fr)
    }
}

impl From<f64> for ContextValue {
    fn from(fr: f64) -> Self {
        ContextValue::Float(fr)
    }
}

impl From<bool> for ContextValue {
    fn from(fr: bool) -> Self {
        ContextValue::Bool(fr)
    }
}

pub struct Context {
    pub cmd: Command,
    pub(crate) config_data: Map,
    pub(crate) env_prefix: Option<String>,
    pub(crate) env_enabled: bool,
}

impl Context {
    pub(crate) fn new(cmd: Command) -> Self {
        Self {
            cmd,
            env_prefix: None,
            env_enabled: true,
            config_data: Map::new(),
        }
    }

    pub(crate) fn set_opt<T: Into<ContextValue>>(&mut self, opt: &Opt, value: T) {
        self.config_data.set(opt.name.clone(), value.into());
    }

    pub(crate) fn set_flag(&mut self, opt: &Opt) {
        self.config_data.set(opt.name.clone(), true.into());
    }

    /// Checks for the existance of a flag
    pub fn flag(&self, name: &str) -> bool {
        self.config_data.get(name).is_some() || self.get_env_value(name).is_some()
    }

    fn get_env_value(&self, name: &str) -> Option<String> {
        if !self.env_enabled {
            return None;
        }
        let name = self
            .env_prefix
            .as_ref()
            .map(|pre| format!("{}_{}", pre, name))
            .unwrap_or_else(|| name.to_string());
        env::var(name).ok()
    }

    /// Returns the value of an option if one exists
    pub fn get_value(&self, name: &str) -> Option<ContextValue> {
        self.config_data
            .get(name)
            .cloned()
            .or_else(|| self.get_env_value(name).map(ContextValue::String))
    }

    /// Returns a string for the field if it is a string
    pub fn get_string(&self, name: &str) -> Option<String> {
        match self.get_value(name) {
            Some(ContextValue::String(v)) => Some(v),
            Some(ContextValue::Integer(v)) => Some(v.to_string()),
            Some(ContextValue::Float(v)) => Some(v.to_string()),
            Some(ContextValue::Bool(v)) => Some(v.to_string()),
            _ => None,
        }
    }

    /// Returns a i64 for the field if it is an integer
    pub fn get_int(&self, name: &str) -> Option<i64> {
        match self.get_value(name) {
            Some(ContextValue::String(v)) => v.parse().ok(),
            Some(ContextValue::Integer(v)) => Some(v),
            Some(ContextValue::Float(v)) => Some(v as i64),
            Some(ContextValue::Bool(true)) => Some(1),
            Some(ContextValue::Bool(false)) => Some(0),
            _ => None,
        }
    }

    /// Returns a f64 for the field if it is a float
    pub fn get_float(&self, name: &str) -> Option<f64> {
        match self.get_value(name) {
            Some(ContextValue::String(v)) => v.parse().ok(),
            Some(ContextValue::Integer(v)) => Some(v as f64),
            Some(ContextValue::Float(v)) => Some(v),
            Some(ContextValue::Bool(true)) => Some(1_f64),
            Some(ContextValue::Bool(false)) => Some(0_f64),
            _ => None,
        }
    }

    /// Returns a f64 for the field if it is a float
    pub fn get_bool(&self, name: &str) -> Option<bool> {
        match self.get_value(name) {
            Some(ContextValue::String(v)) => v.parse().ok(),
            Some(ContextValue::Integer(0)) => Some(false),
            Some(ContextValue::Integer(_)) => Some(true),
            Some(ContextValue::Float(v)) => {
                if v == 0_f64 {
                    Some(false)
                } else {
                    Some(true)
                }
            }
            Some(ContextValue::Bool(v)) => Some(v),
            _ => None,
        }
    }

    /// Can be used to display the automatic help message for the current command.
    pub fn display_help(&self) {
        crate::command::print_command_help(&self.cmd, &[])
    }

    pub(crate) fn load_config_file(&mut self, filename: &str) {
        if std::fs::metadata(filename)
            .ok()
            .map(|f| f.is_file())
            .unwrap_or(false)
        {
            let contents = std::fs::read_to_string(filename).unwrap();
            self.load_config_str(&contents);
        }
    }
}

#[cfg(not(feature = "config"))]
mod noconfig {
    use super::Context;

    impl Context {
        pub(crate) fn load_config_str(&mut self, _filename: &str) {}
    }
}

#[cfg(feature = "config_toml")]
mod toml {
    use super::{Array, Context, ContextValue, Map};

    impl From<toml::Value> for ContextValue {
        fn from(fr: toml::Value) -> Self {
            match fr {
                toml::Value::String(v) => ContextValue::String(v),
                toml::Value::Integer(v) => ContextValue::Integer(v),
                toml::Value::Float(v) => ContextValue::Float(v),
                toml::Value::Boolean(v) => ContextValue::Bool(v),
                toml::Value::Array(v) => ContextValue::Array(Array(
                    v.iter().map(|v| ContextValue::from(v.clone())).collect(),
                )),
                toml::Value::Table(v) => ContextValue::Map(Map(v
                    .iter()
                    .map(|(k, v)| (k.clone(), ContextValue::from(v.clone())))
                    .collect())),
                toml::Value::Datetime(v) => ContextValue::String(v.to_string()),
            }
        }
    }

    impl Context {
        pub(crate) fn load_config_str(&mut self, data: &str) {
            self.config_data = ContextValue::from(toml::from_str::<toml::Value>(data).unwrap())
                .map()
                .clone()
        }
    }
}

#[cfg(feature = "config_json")]
mod json {
    use super::{Array, Context, ContextValue, Map};
    use serde_json::Value;

    impl From<Value> for ContextValue {
        fn from(fr: Value) -> Self {
            match fr {
                Value::String(v) => ContextValue::String(v),
                Value::Number(v) if v.is_i64() => ContextValue::Integer(v.as_i64().unwrap_or(0)),
                Value::Number(v) if v.is_f64() => ContextValue::Float(v.as_f64().unwrap_or(0_f64)),
                Value::Number(v) => ContextValue::Integer(v.as_i64().unwrap_or(0)),
                Value::Array(v) => ContextValue::Array(Array(
                    v.iter().map(|v| ContextValue::from(v.clone())).collect(),
                )),
                Value::Bool(v) => ContextValue::Bool(v),
                Value::Null => ContextValue::String("null".to_string()),
                Value::Object(v) => ContextValue::Map(Map(v
                    .iter()
                    .map(|(k, v)| (k.clone(), ContextValue::from(v.clone())))
                    .collect())),
            }
        }
    }

    impl Context {
        pub(crate) fn load_config_str(&mut self, data: &str) {
            self.config_data =
                ContextValue::from(serde_json::from_str::<Value>(data).unwrap_or_default())
                    .map()
                    .clone()
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::{App, Opt};
    #[test]
    fn test_env() {
        std::env::set_var("thing", "1");
        App::new()
            .opt(Opt::flag("thing").short("-t").long("--t"))
            .handler(|_, ctx, _| {
                assert!(ctx.flag("thing"));
            })
            .run_with(vec![])
            .unwrap();

        std::env::set_var("thing", "one");
        App::new()
            .opt(Opt::scalar("thing").short("-t").long("--t"))
            .handler(|_, ctx, _| {
                assert_eq!(ctx.get_string("thing"), Some("one".into()));
            })
            .run_with(vec![])
            .unwrap();

        std::env::set_var("thing", "one");
        App::new()
            .opt(Opt::scalar("thing").short("t").long("thing"))
            .handler(|_, ctx, _| {
                assert_eq!(ctx.get_string("thing"), Some("1".into()));
            })
            .run_with(vec!["-t".into(), "1".into()])
            .unwrap();
    }

    #[cfg(feature = "config_toml")]
    #[test]
    fn config_load_toml() {
        use super::{Command, Context};
        let cmd = Command::new("xxx");
        let mut ctx = Context::new(cmd);
        ctx.load_config_str(
            r#"
                config_thing = "good"
                num = 15
                
                [nested]
                something = 14
            "#,
        );
        assert_eq!(ctx.get_string("config_thing"), Some("good".into()));
    }

    #[cfg(feature = "config_json")]
    #[test]
    fn config_load_toml() {
        use super::{Command, Context};
        let cmd = Command::new("xxx");
        let mut ctx = Context::new(cmd);
        ctx.load_config_str(
            r#"
                { 
                    "config_thing": "good",
                    "num": 15,
                    "nested": {
                        "something": 14
                    }
                }
            "#,
        );
        assert_eq!(ctx.get_string("config_thing"), Some("good".into()));
    }
}