cli-app-base/src/lib.rs

#[macro_use]
pub extern crate log;

use serde::de::DeserializeOwned;
use serde::Serialize;
use std::collections::HashMap;
use std::fmt::Debug;
use std::fs::File;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::str::FromStr;

// re-export libs appearing in public API
pub use anyhow;
pub use clap;

pub const LOG_LEVELS: [&str; 6] = ["off", "error", "warn", "info", "debug", "trace"];

/// Implement this for the main config struct
pub trait AppConfig: Sized + Serialize + DeserializeOwned + Debug + Default {
    type Init;

    /// Get log level
    fn logging(&self) -> &str;

    /// Print startup banner.
    /// May be overridden to customize or disable it.
    fn print_banner(name: &str, version: &str) {
        eprintln!("{} {}\nRun with -h for help", name, version);
    }

    /// Log messages printed before logging is set up
    /// May be overridden to customize or disable it.
    fn pre_log_println(message: String) {
        eprintln!("{}", message);
    }

    /// Get names of library modules to suppress from log output (limit to warn or error)
    fn logging_suppress_mods(&self) -> Option<Vec<&str>> {
        None
    }

    /// Get log module levels to use (take priority over the main log level)
    fn logging_mod_levels(&self) -> Option<&HashMap<String, String>> {
        None
    }

    /// Add args to later use in the `configure` method.
    fn add_args<'a: 'b, 'b>(clap: clap::App<'a, 'b>) -> clap::App<'a, 'b> {
        // Default impl
        clap
    }

    /// Called when the config file is resolved, e.g. to store it in Config
    /// for later relative path resolution
    fn on_config_file_found(&mut self, _path: &PathBuf) {
        //
    }

    /// Configure the config object using args.
    /// Logging has already been inited and configured.
    fn configure<'a>(self, _clap: &clap::ArgMatches<'a>) -> anyhow::Result<Self::Init>;

    /// Initialize the app
    ///
    /// Use `env!("CARGO_PKG_VERSION")` to get a version string from Cargo.toml
    fn init(name: &str, cfg_file_name: &str, version: &str) -> anyhow::Result<Self::Init> {
        let clap = clap::App::new(name)
            .arg(
                clap::Arg::with_name("config")
                    .short("c")
                    .long("config")
                    .value_name("FILE")
                    .help("Sets a custom config file (JSON5)")
                    .takes_value(true),
            )
            .arg(
                clap::Arg::with_name("dump-config")
                    .long("dump-config")
                    .takes_value(false)
                    .help("Print the loaded config struct"),
            )
            .arg(
                clap::Arg::with_name("default-config")
                    .long("default-config")
                    .takes_value(false)
                    .help("Print the default config JSON for reference (or to be piped to a file)"),
            )
            .arg(
                clap::Arg::with_name("verbose")
                    .short("v")
                    .long("verbose")
                    .multiple(true)
                    .help("Increase logging verbosity (repeat to increase)"),
            )
            .arg(
                clap::Arg::with_name("log-level")
                    .long("log")
                    .takes_value(true)
                    .value_name("LEVEL")
                    .validator(|s| {
                        if LOG_LEVELS.contains(&s.as_str()) {
                            return Ok(());
                        }
                        Err(format!("Bad log level: {}", s))
                    })
                    .help("Set logging verbosity (error,warning,info,debug,trace)"),
            );

        let clap = Self::add_args(clap);

        // this must be done after `add_args` or all hell breaks loose around lifetimes
        let clap = clap.version(version);

        let argv = clap.get_matches();

        if argv.is_present("default-config") {
            println!(
                "{}",
                serde_json::to_string_pretty(&Self::default()).unwrap()
            );
            std::process::exit(0);
        }

        /* Load config */
        let cfg_file = argv.value_of("config").unwrap_or(cfg_file_name);

        Self::print_banner(name, version);

        let filepath = PathBuf::from_str(cfg_file)?;

        let config = if filepath.exists() {
            let mut path = filepath.canonicalize()?;

            if path.is_dir() {
                path = path.join(cfg_file_name);
            }

            if path.exists() {
                Self::pre_log_println(format!("Loading config from: {}", path.display()));

                let buf = read_file(&path)?;

                let mut config: Self = json5::from_str(&buf)?;
                config.on_config_file_found(&path);
                config
            } else {
                Self::pre_log_println(format!(
                    "Config file \"{}\" does not exist, using defaults.",
                    path.display()
                ));

                Self::default()
            }
        } else {
            Self::pre_log_println(format!(
                "Config path \"{}\" does not exist, using defaults.",
                cfg_file
            ));

            Self::default()
        };

        let mut level = config.logging();

        if !LOG_LEVELS.contains(&level) {
            anyhow::bail!("Invalid default log level: {}", level);
        }

        /* env RUST_LOG overrides default if set, but can be changed by CLI args */
        let env_level = option_env!("RUST_LOG").unwrap_or("");
        if !env_level.is_empty() {
            level = env_level;

            if !LOG_LEVELS.contains(&level) {
                anyhow::bail!("Invalid env log level: {}", level);
            }
        }

        /* Explicitly requested level */
        if let Some(l) = argv.value_of("log-level") {
            level = l; // validated by clap
        }

        /* Verbosity increased */
        if argv.is_present("verbose") {
            // bump verbosity if -v's are present
            let pos = LOG_LEVELS.iter().position(|x| x == &level).unwrap();

            level = match LOG_LEVELS
                .iter()
                .nth(pos + argv.occurrences_of("verbose") as usize)
            {
                Some(new_level) => new_level,
                None => "trace",
            };
        }

        let env = env_logger::Env::default().default_filter_or(level);
        let mut builder = env_logger::Builder::from_env(env);

        builder.format_timestamp_millis();

        let mut per_mod = vec![];
        if let Some(mod_levels) = config.logging_mod_levels() {
            for (module, lvl) in mod_levels {
                per_mod.push((module, log::LevelFilter::from_str(lvl)?));
            }
        }

        if let Some(suppress_mods) = config.logging_suppress_mods() {
            // set logging level for suppressed libs
            let mut sup_lvl = log::LevelFilter::Info;
            if level == "warn" {
                sup_lvl = log::LevelFilter::Warn;
            } else if level == "error" {
                sup_lvl = log::LevelFilter::Error;
            }

            'sup: for lib in suppress_mods {
                for (module, _lvl) in per_mod.iter() {
                    if lib.starts_with(module.as_str()) {
                        // avoid suppressing if user set different level for a parent module
                        continue 'sup;
                    }
                }

                builder.filter_module(lib, sup_lvl);
            }
        }

        for (module, lvl) in per_mod.into_iter() {
            builder.filter_module(module, lvl);
        }

        builder.init();

        if argv.is_present("dump-config") {
            debug!("Loaded config: {:#?}", config);
        }

        config.configure(&argv)
    }
}

fn read_file<P: AsRef<Path>>(path: P) -> anyhow::Result<String> {
    let path = path.as_ref();
    let mut file = File::open(path)?;

    let mut buf = String::new();
    file.read_to_string(&mut buf)?;
    Ok(buf)
}