| // Copyright 2023 Google LLC |
| // Use of this source code is governed by a BSD-style license that can be |
| // found in the LICENSE file. |
| |
| //! Debug command traits and helpers for defining commands for integration |
| //! into a debug tool. |
| use log::LevelFilter; |
| use std::str::FromStr; |
| use std::sync::Arc; |
| |
| #[must_use = "if unused the previous log level will immediately be restored"] |
| pub struct ScopedVerbosityGuard { |
| prev_level: LevelFilter, |
| } |
| |
| impl ScopedVerbosityGuard { |
| pub fn new(scoped_level: LevelFilter) -> Self { |
| let prev_level = log::max_level(); |
| log::set_max_level(scoped_level); |
| Self { prev_level } |
| } |
| } |
| |
| impl Drop for ScopedVerbosityGuard { |
| fn drop(&mut self) { |
| log::set_max_level(self.prev_level); |
| } |
| } |
| |
| pub fn parse_level_filter(s: &str) -> Option<LevelFilter> { |
| s.trim().parse().ok() |
| } |
| |
| pub fn extract_verbosity_flags(args: Vec<String>) -> (Option<LevelFilter>, Vec<String>) { |
| let mut level = None; |
| let mut clean_args = Vec::new(); |
| let mut iter = args.into_iter().peekable(); |
| |
| while let Some(arg) = iter.next() { |
| match arg.as_str() { |
| "-v" | "--verbose" => { |
| level = Some(LevelFilter::Debug); |
| } |
| "-vv" | "--trace" => { |
| level = Some(LevelFilter::Trace); |
| } |
| "-q" | "--quiet" => { |
| level = Some(LevelFilter::Warn); |
| } |
| "--log-level" => { |
| if let Some(next) = iter.peek() { |
| if let Some(lvl) = parse_level_filter(next) { |
| level = Some(lvl); |
| iter.next(); |
| continue; |
| } |
| } |
| clean_args.push(arg); |
| } |
| _ if arg.starts_with("--log-level=") => { |
| let val = &arg["--log-level=".len()..]; |
| if let Some(lvl) = parse_level_filter(val) { |
| level = Some(lvl); |
| } else { |
| clean_args.push(arg); |
| } |
| } |
| _ => { |
| clean_args.push(arg); |
| } |
| } |
| } |
| |
| (level, clean_args) |
| } |
| |
| pub fn handle_verbose_command(args: &[String]) { |
| match args.len() { |
| 0 => { |
| println!("Current log level: {}", log::max_level()); |
| } |
| 1 => { |
| if let Some(level) = parse_level_filter(&args[0]) { |
| log::set_max_level(level); |
| println!("Log level set to {}", level); |
| } else { |
| eprintln!( |
| "Invalid log level '{}'. Valid levels: off, error, warn, info, debug, trace", |
| args[0] |
| ); |
| } |
| } |
| _ => { |
| println!("Usage: verbose [off|error|warn|info|debug|trace]"); |
| } |
| } |
| } |
| |
| /// A CommandSet is a set of commands (usually an enum) that each represent an |
| /// action that can be performed. i.e. 'list', 'volume' etc. Each command can |
| /// take zero or more arguments and have zero or more flags. |
| /// Typically an Enum of commands would implement CommandSet trait. |
| pub trait CommandSet: FromStr + ::core::fmt::Display { |
| /// Returns a vector of strings that are the commands supported by this. |
| fn variants() -> Vec<String>; |
| |
| /// Returns a string listing the arguments that this command takes, in <> |
| /// brackets |
| fn arguments(&self) -> &'static str; |
| |
| /// Returns a string displaying the flags that this command supports, in [] |
| /// brackets |
| fn flags(&self) -> &'static str; |
| |
| /// Returns a short description of this command |
| fn desc(&self) -> &'static str; |
| |
| /// Help string for this variant (build from Display, arguments and flags by |
| /// default) |
| fn help_simple(&self) -> String { |
| format!("{self} {} {} -- {}", self.flags(), self.arguments(), self.desc()) |
| } |
| |
| /// Possibly multi-line help string for all variants of this set. |
| fn help_all() -> String { |
| let mut help: String = Self::variants() |
| .into_iter() |
| .filter_map(|s| FromStr::from_str(&s).ok()) |
| .map(|s: Self| format!("{}\n", s)) |
| .collect(); |
| if !help.is_empty() { |
| help.push('\n'); |
| } |
| help.push_str("Global verbosity flags: [-v|--verbose] [-vv|--trace] [-q|--quiet] [--log-level=<level>]\n"); |
| help |
| } |
| } |
| |
| /// A top-level CLI command: either a framework meta-command (e.g. `verbose`) |
| /// or a profile-specific domain command. |
| #[derive(Clone, Copy, Debug, PartialEq, Eq)] |
| pub enum CliCommand<T> { |
| Verbose, |
| Domain(T), |
| } |
| |
| impl<T> From<T> for CliCommand<T> { |
| fn from(cmd: T) -> Self { |
| CliCommand::Domain(cmd) |
| } |
| } |
| |
| impl<T: CommandSet> CommandSet for CliCommand<T> { |
| fn variants() -> Vec<String> { |
| let mut variants = T::variants(); |
| variants.push("verbose".to_string()); |
| variants |
| } |
| |
| fn arguments(&self) -> &'static str { |
| match self { |
| CliCommand::Domain(cmd) => cmd.arguments(), |
| CliCommand::Verbose => "[level] ", |
| } |
| } |
| |
| fn flags(&self) -> &'static str { |
| match self { |
| CliCommand::Domain(cmd) => cmd.flags(), |
| CliCommand::Verbose => "", |
| } |
| } |
| |
| fn desc(&self) -> &'static str { |
| match self { |
| CliCommand::Domain(cmd) => cmd.desc(), |
| CliCommand::Verbose => "Get or set the logging verbosity level", |
| } |
| } |
| } |
| |
| impl<T: ::core::fmt::Display> ::core::fmt::Display for CliCommand<T> { |
| fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result { |
| match self { |
| CliCommand::Domain(cmd) => write!(f, "{cmd}"), |
| CliCommand::Verbose => write!(f, "verbose"), |
| } |
| } |
| } |
| |
| impl<T: FromStr> FromStr for CliCommand<T> { |
| type Err = T::Err; |
| |
| fn from_str(s: &str) -> Result<Self, Self::Err> { |
| if s == "verbose" { |
| Ok(CliCommand::Verbose) |
| } else { |
| T::from_str(s).map(CliCommand::Domain) |
| } |
| } |
| } |
| |
| /// Macro to help build CommandSets |
| #[macro_export] |
| macro_rules! gen_commandset { |
| ($name:ident { |
| $($variant:ident = ($val:expr, [$($flag:expr),*], [$($arg:expr),*], $help:expr)),*, |
| }) => { |
| /// Enum of all possible commands |
| #[derive(PartialEq, Debug, Clone, Copy)] |
| pub enum $name { |
| $($variant,)* |
| } |
| |
| impl CommandSet for $name { |
| fn variants() -> Vec<String> { |
| let mut variants = Vec::new(); |
| $(variants.push($val.to_string());)* |
| variants |
| } |
| |
| fn arguments(&self) -> &'static str { |
| match self { |
| $( |
| $name::$variant => concat!($("<", $arg, "> ",)*), |
| )* |
| } |
| } |
| |
| fn flags(&self) -> &'static str { |
| match self { |
| $( |
| $name::$variant => concat!($("[", $flag, "] ",)*), |
| )* |
| } |
| } |
| |
| fn desc(&self) -> &'static str { |
| match self { |
| $( |
| $name::$variant => $help, |
| )* |
| } |
| } |
| } |
| |
| impl ::core::fmt::Display for $name { |
| fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result { |
| match *self { |
| $($name::$variant => write!(f, $val),)* |
| } |
| } |
| } |
| |
| impl ::std::str::FromStr for $name { |
| type Err = (); |
| |
| fn from_str(s: &str) -> Result<$name, ()> { |
| match s { |
| $($val => Ok($name::$variant),)* |
| _ => Err(()), |
| } |
| } |
| } |
| } |
| } |
| |
| /// CommandRunner is used to perform a specific task based on the command set. |
| pub trait CommandRunner { |
| type Set: CommandSet; |
| type Error: ::std::error::Error; |
| |
| fn run_command( |
| &self, |
| cmd: Self::Set, |
| args: Vec<String>, |
| ) -> impl futures::Future<Output = Result<(), Self::Error>>; |
| |
| fn run( |
| &self, |
| cmd: impl Into<CliCommand<Self::Set>>, |
| args: Vec<String>, |
| ) -> impl futures::Future<Output = Result<(), Self::Error>> { |
| let cmd = cmd.into(); |
| async move { |
| match cmd { |
| CliCommand::Verbose => { |
| handle_verbose_command(&args); |
| Ok(()) |
| } |
| CliCommand::Domain(cmd) => { |
| let (scoped_level, clean_args) = extract_verbosity_flags(args); |
| let _guard = scoped_level.map(ScopedVerbosityGuard::new); |
| self.run_command(cmd, clean_args).await |
| } |
| } |
| } |
| } |
| } |
| |
| impl<R: CommandRunner> CommandRunner for &R { |
| type Set = R::Set; |
| type Error = R::Error; |
| |
| fn run_command( |
| &self, |
| cmd: Self::Set, |
| args: Vec<String>, |
| ) -> impl futures::Future<Output = Result<(), Self::Error>> { |
| (**self).run_command(cmd, args) |
| } |
| } |
| |
| impl<R: CommandRunner> CommandRunner for Arc<R> { |
| type Set = R::Set; |
| type Error = R::Error; |
| |
| fn run_command( |
| &self, |
| cmd: Self::Set, |
| args: Vec<String>, |
| ) -> impl futures::Future<Output = Result<(), Self::Error>> { |
| (**self).run_command(cmd, args) |
| } |
| } |
| |
| #[cfg(test)] |
| mod tests { |
| use super::*; |
| |
| #[test] |
| fn gen_commandset_simple() { |
| gen_commandset! { |
| TestCmd { |
| One = ("one", [], [], "First Command"), |
| WithFlags = ("with-flags", ["-1","-2"], [], "Command with flags"), |
| WithArgs = ("with-args", [], ["arg", "two"], "Command with args"), |
| WithBoth = ("with-both", ["-w"], ["simple"], "Command with both flags and args"), |
| } |
| } |
| |
| let cmd: TestCmd = "one".parse().unwrap(); |
| assert_eq!(cmd, TestCmd::One); |
| |
| let cmd2: TestCmd = "with-flags".parse().unwrap(); |
| assert_eq!(cmd2.arguments(), ""); |
| assert_eq!(cmd2.flags(), "[-1] [-2] "); |
| |
| let verbose_cmd: CliCommand<TestCmd> = "verbose".parse().unwrap(); |
| assert_eq!(verbose_cmd, CliCommand::Verbose); |
| assert_eq!(verbose_cmd.flags(), ""); |
| assert_eq!(verbose_cmd.to_string(), "verbose"); |
| |
| let domain_cmd: CliCommand<TestCmd> = "one".parse().unwrap(); |
| assert_eq!(domain_cmd, CliCommand::Domain(TestCmd::One)); |
| |
| assert!(!TestCmd::variants().contains(&"verbose".to_string())); |
| assert!(CliCommand::<TestCmd>::variants().contains(&"verbose".to_string())); |
| assert!(TestCmd::help_all().contains("Global verbosity flags:")); |
| assert!(CliCommand::<TestCmd>::help_all().contains("Global verbosity flags:")); |
| } |
| |
| #[test] |
| fn test_extract_verbosity_flags() { |
| // -v / --verbose |
| let (lvl, args) = extract_verbosity_flags(vec!["-v".to_string()]); |
| assert_eq!(lvl, Some(LevelFilter::Debug)); |
| assert!(args.is_empty()); |
| |
| let (lvl, args) = extract_verbosity_flags(vec!["--verbose".to_string()]); |
| assert_eq!(lvl, Some(LevelFilter::Debug)); |
| assert!(args.is_empty()); |
| |
| // -vv / --trace |
| let (lvl, args) = extract_verbosity_flags(vec!["-vv".to_string()]); |
| assert_eq!(lvl, Some(LevelFilter::Trace)); |
| assert!(args.is_empty()); |
| |
| let (lvl, args) = extract_verbosity_flags(vec!["--trace".to_string()]); |
| assert_eq!(lvl, Some(LevelFilter::Trace)); |
| assert!(args.is_empty()); |
| |
| // -q / --quiet |
| let (lvl, args) = extract_verbosity_flags(vec!["-q".to_string()]); |
| assert_eq!(lvl, Some(LevelFilter::Warn)); |
| assert!(args.is_empty()); |
| |
| let (lvl, args) = extract_verbosity_flags(vec!["--quiet".to_string()]); |
| assert_eq!(lvl, Some(LevelFilter::Warn)); |
| assert!(args.is_empty()); |
| |
| // --log-level=<level> and --log-level <level> (case-insensitive) |
| for (name, expected) in [ |
| ("off", LevelFilter::Off), |
| ("error", LevelFilter::Error), |
| ("warn", LevelFilter::Warn), |
| ("info", LevelFilter::Info), |
| ("debug", LevelFilter::Debug), |
| ("trace", LevelFilter::Trace), |
| ("DEBUG", LevelFilter::Debug), |
| ("Info", LevelFilter::Info), |
| ] { |
| let (lvl, args) = extract_verbosity_flags(vec![format!("--log-level={name}")]); |
| assert_eq!(lvl, Some(expected)); |
| assert!(args.is_empty()); |
| |
| let (lvl, args) = |
| extract_verbosity_flags(vec!["--log-level".to_string(), name.to_string()]); |
| assert_eq!(lvl, Some(expected)); |
| assert!(args.is_empty()); |
| } |
| |
| // Mixed arguments |
| let input = vec![ |
| "action".to_string(), |
| "-v".to_string(), |
| "--custom-flag".to_string(), |
| "target".to_string(), |
| "--log-level=warn".to_string(), |
| "extra".to_string(), |
| ]; |
| let (lvl, args) = extract_verbosity_flags(input); |
| assert_eq!(lvl, Some(LevelFilter::Warn)); |
| assert_eq!(args, vec!["action", "--custom-flag", "target", "extra"]); |
| |
| // Invalid log level flag should not be extracted |
| let (lvl, args) = |
| extract_verbosity_flags(vec!["--log-level=invalid".to_string(), "arg".to_string()]); |
| assert_eq!(lvl, None); |
| assert_eq!(args, vec!["--log-level=invalid", "arg"]); |
| } |
| |
| // Mutex to avoid parallel tests fighting over log levels |
| static LOG_MUTEX: std::sync::Mutex<()> = std::sync::Mutex::new(()); |
| |
| #[test] |
| fn test_scoped_verbosity_guard() { |
| let _lock = LOG_MUTEX.lock().unwrap(); |
| log::set_max_level(LevelFilter::Info); |
| |
| futures::executor::block_on(async { |
| assert_eq!(log::max_level(), LevelFilter::Info); |
| { |
| let _guard = ScopedVerbosityGuard::new(LevelFilter::Trace); |
| assert_eq!(log::max_level(), LevelFilter::Trace); |
| // Yield across async boundary |
| futures::future::ready(()).await; |
| assert_eq!(log::max_level(), LevelFilter::Trace); |
| } |
| assert_eq!(log::max_level(), LevelFilter::Info); |
| }); |
| } |
| |
| #[test] |
| fn test_persistent_verbosity_and_handle_command() { |
| let _lock = LOG_MUTEX.lock().unwrap(); |
| log::set_max_level(LevelFilter::Debug); |
| |
| // handle_verbose_command with 0 args |
| handle_verbose_command(&[]); |
| assert_eq!(log::max_level(), LevelFilter::Debug); |
| |
| // handle_verbose_command with 1 valid arg |
| handle_verbose_command(&["trace".to_string()]); |
| assert_eq!(log::max_level(), LevelFilter::Trace); |
| |
| handle_verbose_command(&["WARN".to_string()]); |
| assert_eq!(log::max_level(), LevelFilter::Warn); |
| |
| // handle_verbose_command with invalid arg (should not change) |
| handle_verbose_command(&["invalid_level".to_string()]); |
| assert_eq!(log::max_level(), LevelFilter::Warn); |
| |
| // handle_verbose_command with > 1 args (should not change) |
| handle_verbose_command(&["info".to_string(), "extra".to_string()]); |
| assert_eq!(log::max_level(), LevelFilter::Warn); |
| |
| // Reset back to Info |
| log::set_max_level(LevelFilter::Info); |
| } |
| |
| #[test] |
| fn test_command_runner_run() { |
| gen_commandset! { |
| RunnerCmd { |
| Action = ("action", [], ["target"], "Perform action"), |
| } |
| } |
| |
| struct TestRunner { |
| executed: std::sync::Mutex<Vec<(RunnerCmd, Vec<String>)>>, |
| } |
| |
| impl CommandRunner for TestRunner { |
| type Set = RunnerCmd; |
| type Error = std::io::Error; |
| |
| fn run_command( |
| &self, |
| cmd: Self::Set, |
| args: Vec<String>, |
| ) -> impl futures::Future<Output = Result<(), Self::Error>> { |
| self.executed.lock().unwrap().push((cmd, args)); |
| futures::future::ready(Ok(())) |
| } |
| } |
| |
| let runner = TestRunner { executed: std::sync::Mutex::new(Vec::new()) }; |
| let _lock = LOG_MUTEX.lock().unwrap(); |
| |
| // Running Verbose command should update persistent verbosity and NOT |
| // invoke run_command |
| log::set_max_level(LevelFilter::Info); |
| futures::executor::block_on(runner.run(CliCommand::Verbose, vec!["trace".to_string()])) |
| .unwrap(); |
| assert_eq!(log::max_level(), LevelFilter::Trace); |
| assert!(runner.executed.lock().unwrap().is_empty()); |
| |
| // Running normal command with verbosity flag strips flag and runs |
| // command |
| let args = vec!["-v".to_string(), "my_target".to_string()]; |
| futures::executor::block_on(runner.run(RunnerCmd::Action, args)).unwrap(); |
| let executed = runner.executed.lock().unwrap().clone(); |
| assert_eq!(executed.len(), 1); |
| assert_eq!(executed[0], (RunnerCmd::Action, vec!["my_target".to_string()])); |
| |
| // After run completes, the guard has restored the previous level |
| assert_eq!(log::max_level(), LevelFilter::Trace); |
| |
| // Reset persistent verbosity back to Info |
| log::set_max_level(LevelFilter::Info); |
| } |
| } |