| // Copyright 2019 The Chromium OS Authors. All rights reserved. |
| // Use of this source code is governed by a BSD-style license that can be |
| // found in the LICENSE file. |
| |
| // Provides support for registering and handling commands as well as displaying the command line |
| // help. |
| |
| use std::collections::HashMap; |
| use std::fmt::{self, Display}; |
| use std::io::Write; |
| use std::process::Child; |
| |
| use remain::sorted; |
| |
| const INDENT: &str = " "; |
| |
| #[sorted] |
| pub enum Error { |
| CommandInvalidArguments(String), |
| CommandNotFound(String), |
| CommandNotImplemented(String), |
| CommandReturnedError, |
| DuplicateCommand(Vec<String>), |
| FlagFilter, |
| } |
| |
| impl Display for Error { |
| #[remain::check] |
| fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { |
| use self::Error::*; |
| |
| #[sorted] |
| match self { |
| CommandInvalidArguments(msg) => write!(f, "invalid arguments: {}", msg), |
| CommandNotFound(command) => write!(f, "unknown command: {}", command), |
| CommandNotImplemented(command) => write!(f, "command not implemented: {}", command), |
| CommandReturnedError => write!(f, "command failed"), |
| DuplicateCommand(dups) => write!(f, "duplicate commands: {}", dups.join(", ")), |
| FlagFilter => write!(f, "error filtering flags"), |
| } |
| } |
| } |
| |
| pub fn wait_for_result(mut child: Child) -> Result<(), Error> { |
| match child.wait() { |
| Ok(status) => { |
| if status.success() { |
| Ok(()) |
| } else { |
| Err(Error::CommandReturnedError) |
| } |
| } |
| Err(_) => Err(Error::CommandReturnedError), |
| } |
| } |
| |
| // Keeps all the state required to interpret a dispatched command, and provides interfaces for: |
| // * Registering commands (along with strings required to generate help text). |
| // * Dispatching an issued command. |
| // * Providing token completion given partial input. |
| pub struct Dispatcher { |
| registered_commands: Vec<Command>, |
| } |
| |
| impl Dispatcher { |
| pub fn new() -> Dispatcher { |
| Dispatcher { |
| registered_commands: Vec::new(), |
| } |
| } |
| |
| // Register a command description that can be handled by the dispatcher. |
| pub fn register_command(&mut self, cmd: Command) -> &mut Dispatcher { |
| self.registered_commands.push(cmd); |
| self |
| } |
| |
| // Lookup a command by name. |
| pub fn find_by_name(&self, name: &str) -> Option<&Command> { |
| find_by_name(name, &self.registered_commands) |
| } |
| |
| // Return a CompletionResult that represents auto-completion suggestions for |tokens|. |
| pub fn complete_command(&self, tokens: Vec<String>) -> CompletionResult { |
| if tokens.is_empty() { |
| return complete_by_name("", &self.registered_commands); |
| } |
| |
| let (commands, entry) = self.get_command_list(tokens); |
| |
| if commands.is_empty() { |
| if entry.tokens.len() == 1 { |
| return complete_by_name(&entry.tokens[0], &self.registered_commands); |
| } |
| return CompletionResult::NoMatches; |
| } |
| |
| let command: &Command = commands.last().unwrap(); |
| if let Some(cb) = command.completion_callback { |
| return (cb)(&entry); |
| } |
| |
| CompletionResult::NoMatches |
| } |
| |
| // Execute the command handler represented by |tokens|. Flags will be parsed and flag handling |
| // callbacks will be invoked. |
| pub fn handle_command(&self, tokens: Vec<String>) -> Result<(), Error> { |
| if tokens.is_empty() { |
| return Err(Error::CommandNotFound(tokens.join(" "))); |
| } |
| |
| let mut command: &Command = self |
| .find_by_name(&tokens[0]) |
| .ok_or_else(|| Error::CommandNotFound(tokens[0].to_string()))?; |
| |
| let mut flag_callbacks: Vec<CommandCallback> = Vec::new(); |
| let entry = &mut Arguments { |
| tokens, |
| position: 1, |
| flags: HashMap::new(), |
| }; |
| |
| if let Some(cb) = command.flag_callback { |
| flag_callbacks.push(cb); |
| } |
| while entry.position < entry.tokens.len() { |
| let sub: Option<&Command> = command.handle_tokens(entry); |
| |
| if sub.is_none() { |
| break; |
| } |
| entry.position += 1; |
| command = sub.unwrap(); |
| if let Some(cb) = command.flag_callback { |
| flag_callbacks.push(cb); |
| } |
| } |
| if command.command_callback.is_none() { |
| return Err(Error::CommandNotImplemented(entry.get_command().join(" "))); |
| } |
| |
| for cb in flag_callbacks { |
| (cb)(&command, entry)?; |
| } |
| (command.command_callback.unwrap())(&command, entry) |
| } |
| |
| pub fn validate(&mut self) -> Result<(), Error> { |
| self.registered_commands |
| .sort_unstable_by(|a: &Command, b: &Command| a.name.cmp(&b.name)); |
| |
| let mut duplicates: Vec<String> = Vec::new(); |
| for i in 1..self.registered_commands.len() { |
| let name = &self.registered_commands[i - 1].name; |
| if name == &self.registered_commands[i].name { |
| duplicates.push(name.to_string()); |
| } |
| } |
| |
| if !duplicates.is_empty() { |
| return Err(Error::DuplicateCommand(duplicates)); |
| } |
| Ok(()) |
| } |
| |
| // Generate and return the help string. |
| pub fn help_string(&self, w: &mut dyn Write, opt_cmds: Option<&[&str]>) -> Result<(), Error> { |
| match opt_cmds { |
| Some(cmds) => { |
| for name in cmds { |
| if let Some(cmd) = self.find_by_name(name) { |
| cmd.append_help_string(w, 0); |
| } else { |
| return Err(Error::CommandNotFound(name.to_string())); |
| } |
| } |
| } |
| None => { |
| for c in &self.registered_commands { |
| c.append_help_string(w, 0); |
| } |
| } |
| } |
| Ok(()) |
| } |
| |
| fn get_command_list(&self, tokens: Vec<String>) -> (Vec<&Command>, Arguments) { |
| let mut list: Vec<&Command> = Vec::new(); |
| let mut entry = Arguments::new(); |
| if tokens.is_empty() { |
| return (list, entry); |
| } |
| entry.tokens = tokens; |
| |
| let c = self.find_by_name(&entry.tokens[0]); |
| if c.is_none() { |
| return (list, entry); |
| } |
| |
| entry.position = 1; |
| let mut command: &Command = c.unwrap(); |
| list.push(command); |
| while entry.position < entry.tokens.len() { |
| let sub: Option<&Command> = command.handle_tokens(&mut entry); |
| |
| if sub.is_none() { |
| break; |
| } |
| entry.position += 1; |
| command = sub.unwrap(); |
| list.push(command); |
| } |
| (list, entry) |
| } |
| } |
| |
| // Owns the data required to identify the command, and serves as a node in a tree. Sub commands can |
| // be registered. It contains callbacks for the following: |
| // * flag/switch processing |
| // * executing the command |
| // * providing command completion suggestions. |
| pub struct Command { |
| name: String, |
| usage: String, |
| description: String, |
| sub_commands: Vec<Command>, |
| flags: Vec<Flag>, |
| flag_callback: Option<CommandCallback>, |
| command_callback: Option<CommandCallback>, |
| completion_callback: Option<CompletionCallback>, |
| help_callback: HelpCallback, |
| } |
| |
| impl Command { |
| pub fn new(name: String, usage: String, description: String) -> Command { |
| Command { |
| name, |
| usage, |
| description, |
| sub_commands: Vec::new(), |
| flags: Vec::new(), |
| flag_callback: None, |
| command_callback: None, |
| completion_callback: None, |
| help_callback: default_help_callback, |
| } |
| } |
| |
| // Set the callback that is executed when this command or a sub command is invoked primarily for |
| // the purpose of filtering or handling flags. |
| pub fn set_flag_callback(mut self, replacement: Option<CommandCallback>) -> Command { |
| self.flag_callback = replacement; |
| self |
| } |
| |
| // Set the callback that is executed when this command invoked directly. |
| pub fn set_command_callback(mut self, replacement: Option<CommandCallback>) -> Command { |
| self.command_callback = replacement; |
| self |
| } |
| |
| // Set the callback to handle command completion for arguments of this command. |
| pub fn set_completion_callback(mut self, replacement: Option<CompletionCallback>) -> Command { |
| self.completion_callback = replacement; |
| self |
| } |
| |
| // Set the callback to handle command completion for arguments of this command. |
| pub fn set_help_callback(mut self, replacement: HelpCallback) -> Command { |
| self.help_callback = replacement; |
| self |
| } |
| |
| pub fn register_subcommand(mut self, cmd: Command) -> Command { |
| self.sub_commands.push(cmd); |
| self |
| } |
| |
| pub fn register_flag(mut self, flag: Flag) -> Command { |
| self.flags.push(flag); |
| self |
| } |
| |
| pub fn get_name(&self) -> &str { |
| &self.name |
| } |
| |
| pub fn append_help_string(&self, w: &mut dyn Write, level: usize) { |
| (self.help_callback)(self, w, level); |
| } |
| |
| fn find_flag(&self, flag: &str) -> Option<&Flag> { |
| for f in &self.flags { |
| if f.name == flag { |
| return Some(&f); |
| } |
| } |
| None |
| } |
| |
| fn find_subcommand(&self, name: &str) -> Option<&Command> { |
| find_by_name(name, &self.sub_commands) |
| } |
| |
| fn handle_tokens(&self, entry: &mut Arguments) -> Option<&Command> { |
| while entry.position < entry.tokens.len() { |
| let token = &entry.tokens[entry.position]; |
| |
| let result = self.find_subcommand(token); |
| if result.is_some() { |
| return result; |
| } |
| |
| let mut parts = token.splitn(2, '='); |
| let name = parts.next().unwrap().to_string(); |
| let flag = self.find_flag(&name)?; |
| |
| let value = flag.get_default_value().parse(parts.next().unwrap_or("")); |
| if value.is_none() { |
| panic!(); |
| } |
| entry.flags.insert(name, value.unwrap()); |
| |
| entry.position += 1; |
| } |
| None |
| } |
| } |
| |
| pub fn default_help_callback(cmd: &Command, w: &mut dyn Write, level: usize) { |
| let mut prefix = INDENT.repeat(level); |
| write!(w, "{}{}", &prefix, &cmd.name).unwrap(); |
| prefix.push_str(&INDENT); |
| if !cmd.usage.is_empty() { |
| write!(w, " {}", &cmd.usage).unwrap(); |
| } |
| writeln!(w).unwrap(); |
| |
| if !cmd.description.is_empty() { |
| writeln!(w, "{}{}", &prefix, &cmd.description).unwrap(); |
| } |
| |
| if !cmd.flags.is_empty() { |
| if !cmd.description.is_empty() { |
| writeln!(w).unwrap(); |
| } |
| write!(w, "{}Options:", &prefix).unwrap(); |
| for flag in &cmd.flags { |
| writeln!(w).unwrap(); |
| flag.append_help_string(w, level + 2); |
| } |
| } |
| |
| if !cmd.sub_commands.is_empty() { |
| if !cmd.description.is_empty() || !cmd.flags.is_empty() { |
| writeln!(w).unwrap(); |
| } |
| writeln!(w, "{}Subcommands:", &prefix).unwrap(); |
| for sub in &cmd.sub_commands { |
| sub.append_help_string(w, level + 2); |
| } |
| } else { |
| writeln!(w).unwrap(); |
| } |
| } |
| |
| impl HasName for Command { |
| fn get_name(&self) -> &str { |
| &self.name |
| } |
| } |
| |
| pub struct Arguments { |
| tokens: Vec<String>, |
| position: usize, |
| flags: HashMap<String, FlagType>, |
| } |
| |
| impl Arguments { |
| fn new() -> Arguments { |
| Arguments { |
| tokens: Vec::new(), |
| position: 0, |
| flags: HashMap::new(), |
| } |
| } |
| |
| pub fn get_command(&self) -> &[String] { |
| &self.tokens[0..self.position] |
| } |
| |
| pub fn get_tokens(&self) -> &[String] { |
| &self.tokens |
| } |
| |
| pub fn get_flag(&self, flag: &str) -> Option<&FlagType> { |
| self.flags.get(flag) |
| } |
| |
| pub fn get_args(&self) -> &[String] { |
| &self.tokens[self.position..] |
| } |
| } |
| |
| pub enum CompletionResult { |
| NoMatches, |
| SingleDiff(String), |
| WholeTokenList(Vec<String>), |
| } |
| |
| type CommandCallback = fn(cmd: &Command, args: &Arguments) -> Result<(), Error>; |
| type CompletionCallback = fn(args: &Arguments) -> CompletionResult; |
| type HelpCallback = fn(cmd: &Command, w: &mut dyn Write, level: usize); |
| |
| pub enum FlagType { |
| NoValue, |
| Boolean(bool), |
| Integer(i64), |
| Float(f64), |
| String(String), |
| } |
| |
| impl FlagType { |
| pub fn parse(&self, s: &str) -> Option<FlagType> { |
| match self { |
| FlagType::NoValue => Some(FlagType::NoValue), |
| FlagType::Boolean(_) => match s.parse::<bool>() { |
| Ok(b) => Some(FlagType::Boolean(b)), |
| Err(_) => None, |
| }, |
| FlagType::Integer(_) => match s.parse::<i64>() { |
| Ok(i) => Some(FlagType::Integer(i)), |
| Err(_) => None, |
| }, |
| FlagType::Float(_) => match s.parse::<f64>() { |
| Ok(f) => Some(FlagType::Float(f)), |
| Err(_) => None, |
| }, |
| FlagType::String(_) => Some(FlagType::String(s.to_string())), |
| } |
| } |
| } |
| |
| pub struct Flag { |
| name: String, |
| description: String, |
| default_value: FlagType, |
| } |
| |
| impl HasName for Flag { |
| fn get_name(&self) -> &str { |
| &self.name |
| } |
| } |
| |
| impl Flag { |
| pub fn new(name: String, description: String, default_value: FlagType) -> Flag { |
| Flag { |
| name, |
| description, |
| default_value, |
| } |
| } |
| |
| pub fn append_help_string(&self, w: &mut dyn Write, level: usize) { |
| let prefix = INDENT.repeat(level); |
| write!(w, "{}{}", &prefix, &self.name).unwrap(); |
| let value_cb = |w: &mut dyn Write, value: String| { |
| write!(w, "{}{}{}", &value, &prefix, INDENT).unwrap(); |
| }; |
| match &self.default_value { |
| FlagType::NoValue => { |
| write!(w, " ").unwrap(); |
| } |
| FlagType::Boolean(default) => { |
| value_cb(w, format!("=[true|false] default: {}\n", default)); |
| } |
| FlagType::Integer(default) => { |
| value_cb(w, format!("=<int> default: {}\n", default)); |
| } |
| FlagType::Float(default) => { |
| value_cb(w, format!("=<float> default: {}\n", default)); |
| } |
| FlagType::String(default) => { |
| value_cb(w, format!("=<value> default: {}\n", default)); |
| } |
| } |
| writeln!(w, "{}", &self.description).unwrap(); |
| } |
| |
| pub fn get_default_value(&self) -> &FlagType { |
| &self.default_value |
| } |
| } |
| |
| // Internal trait used to share logic used to walk lists of Commands and Flags. |
| trait HasName { |
| fn get_name(&self) -> &str; |
| } |
| |
| // Fetch a reference to an entry in a list by matching against get_name(). |
| fn find_by_name<'a, T: HasName>(name: &str, list: &'a [T]) -> Option<&'a T> { |
| for c in list { |
| if *c.get_name() == *name { |
| return Some(c); |
| } |
| } |
| None |
| } |
| |
| // Provide a CompletionResult after prefix matching against get_name(). |
| fn complete_by_name<T: HasName>(name: &str, list: &[T]) -> CompletionResult { |
| let mut suggestions: Vec<String> = Vec::new(); |
| for c in list { |
| if c.get_name().starts_with(name) { |
| suggestions.push(c.get_name().to_string()); |
| } |
| } |
| if suggestions.is_empty() { |
| CompletionResult::NoMatches |
| } else if suggestions.len() == 1 { |
| CompletionResult::SingleDiff(suggestions[0][name.len()..].to_string()) |
| } else { |
| CompletionResult::WholeTokenList(suggestions) |
| } |
| } |
| |
| #[cfg(test)] |
| mod tests { |
| use super::*; |
| |
| use std::str; |
| |
| static PARENT_COMMAND_NAME: &str = "test"; |
| static CHILD_COMMAND_NAME: &str = "subtest"; |
| |
| fn flag_callback(_cmd: &Command, _args: &Arguments) -> Result<(), Error> { |
| Ok(()) |
| } |
| |
| fn false_flag_callback(_cmd: &Command, _args: &Arguments) -> Result<(), Error> { |
| Err(Error::FlagFilter) |
| } |
| |
| fn panic_flag_callback(_cmd: &Command, _args: &Arguments) -> Result<(), Error> { |
| panic!() |
| } |
| |
| fn command_callback(_cmd: &Command, _args: &Arguments) -> Result<(), Error> { |
| Ok(()) |
| } |
| |
| fn false_command_callback(_cmd: &Command, _args: &Arguments) -> Result<(), Error> { |
| Err(Error::CommandReturnedError) |
| } |
| |
| fn panic_command_callback(_cmd: &Command, _args: &Arguments) -> Result<(), Error> { |
| panic!() |
| } |
| |
| fn completion_callback(_args: &Arguments) -> CompletionResult { |
| CompletionResult::NoMatches |
| } |
| |
| fn panic_completion_callback(_args: &Arguments) -> CompletionResult { |
| panic!() |
| } |
| |
| fn default_dispatcher(parent: Command) -> Dispatcher { |
| let mut dispatcher = Dispatcher::new(); |
| dispatcher.register_command(parent); |
| dispatcher |
| } |
| |
| fn default_command(name: String, usage: String, description: String) -> Command { |
| Command::new(name, usage, description) |
| .set_flag_callback(Some(panic_flag_callback)) |
| .set_command_callback(Some(panic_command_callback)) |
| .set_completion_callback(Some(panic_completion_callback)) |
| } |
| |
| fn default_parent_command(child: Command) -> Command { |
| default_command( |
| PARENT_COMMAND_NAME.to_string(), |
| format!("[{}]", CHILD_COMMAND_NAME), |
| "parent test command.".to_string(), |
| ) |
| .register_subcommand(child) |
| } |
| |
| fn default_child_command() -> Command { |
| default_command( |
| CHILD_COMMAND_NAME.to_string(), |
| "".to_string(), |
| "parent test command.".to_string(), |
| ) |
| } |
| |
| #[test] |
| fn test_handle_command_empty() { |
| let dispatcher = default_dispatcher(default_parent_command(default_child_command())); |
| |
| let mut tokens: Vec<String> = Vec::new(); |
| tokens.push(PARENT_COMMAND_NAME.to_string()); |
| |
| assert!(!dispatcher.handle_command(Vec::new()).is_ok()); |
| } |
| |
| #[test] |
| fn test_handle_command_parent() { |
| let dispatcher = default_dispatcher( |
| default_parent_command(default_child_command().set_flag_callback(Some(flag_callback))) |
| .set_flag_callback(Some(flag_callback)) |
| .set_command_callback(Some(command_callback)), |
| ); |
| |
| let mut tokens: Vec<String> = Vec::new(); |
| tokens.push(PARENT_COMMAND_NAME.to_string()); |
| |
| assert!(dispatcher.handle_command(tokens).is_ok()); |
| } |
| |
| #[test] |
| fn test_handle_command_child() { |
| let dispatcher = default_dispatcher( |
| default_parent_command( |
| default_child_command() |
| .set_flag_callback(Some(flag_callback)) |
| .set_command_callback(Some(command_callback)), |
| ) |
| .set_flag_callback(Some(flag_callback)), |
| ); |
| |
| let mut tokens: Vec<String> = Vec::new(); |
| tokens.push(PARENT_COMMAND_NAME.to_string()); |
| tokens.push(CHILD_COMMAND_NAME.to_string()); |
| |
| assert!(dispatcher.handle_command(tokens).is_ok()); |
| } |
| |
| #[test] |
| fn test_handle_command_false_flag_callback() { |
| let dispatcher = default_dispatcher( |
| default_parent_command(default_child_command()) |
| .set_flag_callback(Some(false_flag_callback)), |
| ); |
| |
| let mut tokens: Vec<String> = Vec::new(); |
| tokens.push(PARENT_COMMAND_NAME.to_string()); |
| tokens.push(CHILD_COMMAND_NAME.to_string()); |
| |
| assert!(!dispatcher.handle_command(tokens).is_ok()); |
| } |
| |
| #[test] |
| fn test_help_string() { |
| let dispatcher = default_dispatcher( |
| default_command("1".to_string(), "2".to_string(), "3".to_string()) |
| .register_subcommand(default_command( |
| "4".to_string(), |
| "5".to_string(), |
| "6".to_string(), |
| )) |
| .register_subcommand( |
| default_command("7".to_string(), "".to_string(), "".to_string()) |
| .register_subcommand( |
| default_command("8".to_string(), "9".to_string(), "10".to_string()) |
| .register_flag(Flag::new( |
| "11".to_string(), |
| "12".to_string(), |
| FlagType::Integer(0), |
| )), |
| ), |
| ), |
| ); |
| |
| let mut result: Vec<u8> = Vec::new(); |
| assert!(dispatcher |
| .help_string(&mut result, Some(&["not a command"])) |
| .is_err()); |
| assert!(dispatcher.help_string(&mut result, None).is_ok()); |
| assert_eq!( |
| str::from_utf8(&result).unwrap(), |
| r#"1 2 |
| 3 |
| |
| Subcommands: |
| 4 5 |
| 6 |
| |
| 7 |
| Subcommands: |
| 8 9 |
| 10 |
| |
| Options: |
| 11=<int> default: 0 |
| 12 |
| |
| "# |
| ); |
| } |
| } |