From 81e557a037b6e037e5c1d703db3696c8525eeb5b Mon Sep 17 00:00:00 2001 From: Joe Bellus Date: Sat, 19 Jun 2021 15:51:11 -0400 Subject: [PATCH] Environment based options/flags Environment variables can be checked for options/flags. The name of the option is used to query if an environment variable exists for the given option. This can optionally be prefixed. --- README.md | 1 + benches/arg_parsing.rs | 4 +-- examples/fib.rs | 2 +- src/app.rs | 59 ++++++++++++++++++++++++++++++++++++----- src/context.rs | 60 ++++++++++++++++++++++++++++++++++++++++-- 5 files changed, 114 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 331e55f..5146c26 100644 --- a/README.md +++ b/README.md @@ -10,6 +10,7 @@ building blocks for building attractive and smooth CLIs * Nested subcommands with their own flags * Opts are hierarchal and can be utilized from parent commands * Automatic usage details for subcommands and bare execution +* Flags and options can be set via environment variables ## Styling diff --git a/benches/arg_parsing.rs b/benches/arg_parsing.rs index e12573a..c700c5c 100644 --- a/benches/arg_parsing.rs +++ b/benches/arg_parsing.rs @@ -16,8 +16,8 @@ fn parse_args(c: &mut Criterion) { Command::new("thing") .opt(Opt::scalar("config").short("c").long("config")) .handler(|_, ctx, _| { - assert_eq!(ctx.get_string("user"), Some("joe".into())); - assert_eq!(ctx.get_string("config"), Some("c.json".into())); + assert_eq!(ctx.get_value("user"), Some("joe".into())); + assert_eq!(ctx.get_value("config"), Some("c.json".into())); }), ); diff --git a/examples/fib.rs b/examples/fib.rs index f0237ff..25dacd2 100644 --- a/examples/fib.rs +++ b/examples/fib.rs @@ -21,7 +21,7 @@ fn main() { fn fibonacci_handler(_app: &App, ctx: &Context, _args: &[String]) { let v = fibonacci( - ctx.get_string("count") + ctx.get_value("count") .unwrap_or("1".to_string()) .parse() .unwrap_or(1), diff --git a/src/app.rs b/src/app.rs index 333c9eb..b541c2a 100644 --- a/src/app.rs +++ b/src/app.rs @@ -10,6 +10,7 @@ pub struct App { name: Option<&'static str>, version: Option<&'static str>, pub root: Command, + env_prefix: Option<&'static str>, } impl Default for App { @@ -18,6 +19,7 @@ impl Default for App { name: None, version: None, root: Command::new("root"), + env_prefix: None, } } } @@ -68,6 +70,21 @@ impl App { self } + /// Sets the environment variable prefix for option resolution. If set to something like + /// APP_NAME a option with the name THING, will look for an environment variable named + /// APP_NAME_THING. + /// + /// Example: + /// ```rust + /// use arkham::App; + /// App::new().env_prefix("APP_NAME"); + /// + /// ``` + pub fn env_prefix(mut self, prefix: &'static str) -> Self { + self.env_prefix = Some(prefix); + self + } + /// Adds a root level command to the application. This command can then be executed with: /// /// myapp command_name @@ -131,7 +148,7 @@ impl App { /// .run_with(vec!["-n".to_string(), "alice".to_string()]); /// /// fn my_handler(app: &App, ctx: &Context, args: &[String]) { - /// println!("Hello, {}", ctx.get_string("name").unwrap()); + /// println!("Hello, {}", ctx.get_value("name").unwrap()); /// } /// ``` pub fn run_with(&self, args: Vec) -> Result<()> { @@ -151,7 +168,7 @@ impl App { /// .run(); /// /// fn my_handler(app: &App, ctx: &Context, args: &[String]) { - /// println!("Hello, {}", ctx.get_string("name").unwrap_or_else(|| "unnamed".into())); + /// println!("Hello, {}", ctx.get_value("name").unwrap_or_else(|| "unnamed".into())); /// } /// ``` pub fn run(&mut self) -> Result<()> { @@ -193,7 +210,18 @@ fn run_command(app: &App, cmd: &Command, args: &[String], opts: &mut Vec { + opts.push(ActiveOpt::new(opt.clone(), vec!["".into()])); + } + OptKind::String => { + if let Some(value) = args.next() { + opts.push(ActiveOpt::new(opt.clone(), vec![value.clone()])); + } else { + return Err(OptError::InvalidOpt(opt.name.clone())); + } + } + } } else { ignored.push(arg.clone()); } @@ -225,7 +253,11 @@ fn run_command(app: &App, cmd: &Command, args: &[String], opts: &mut Vec, cmd: Command, + pub(crate) env_prefix: Option, + pub(crate) env_enabled: bool, } impl Context { pub(crate) fn new(cmd: Command, opts: Vec) -> Self { - Self { opts, cmd } + Self { + opts, + cmd, + env_prefix: None, + env_enabled: true, + } } /// Checks for the existance of a flag @@ -18,10 +26,23 @@ impl Context { self.opts .iter() .any(|o| o.definition.name == name && matches!(o.definition.kind, OptKind::Flag)) + || self.get_env_value(name).is_some() + } + + fn get_env_value(&self, name: &str) -> Option { + if self.env_enabled == false { + return None; + } + let name = self + .env_prefix + .as_ref() + .map(|pre| format!("{}_{}", pre, name)) + .unwrap_or(name.to_string()); + env::var(name).ok() } /// Returns the value of an option if one exists - pub fn get_string(&self, name: &str) -> Option { + pub fn get_value(&self, name: &str) -> Option { self.opts .iter() .find_map(|o| { @@ -32,6 +53,7 @@ impl Context { } }) .flatten() + .or(self.get_env_value(name)) } /// Can be used to display the automatic help message for the current command. @@ -39,3 +61,37 @@ impl Context { crate::command::print_command_help(&self.cmd, &vec![]) } } + +#[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_eq!(ctx.flag("thing"), true); + }) + .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_value("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_value("thing"), Some("1".into())); + }) + .run_with(vec!["-t".into(), "1".into()]) + .unwrap(); + } +}