# HG changeset patch # User Arseniy Alekseyev # Date 2022-04-12 19:01:49 # Node ID 86c49b0000c4b999795a0c188ab1bc69d0d8033f # Parent cb21c9c1a48035fb3c9927e203bc6d1aac9b84d7 rhg: refactor to pass argv down, instead of caling args_os() This refactoring makes it easy to patch some command-line preprocessing into rhg. We use this to support using rhg as a shebang interpreter, for example. Differential Revision: https://phab.mercurial-scm.org/D12543 diff --git a/rust/rhg/src/blackbox.rs b/rust/rhg/src/blackbox.rs --- a/rust/rhg/src/blackbox.rs +++ b/rust/rhg/src/blackbox.rs @@ -5,6 +5,7 @@ use format_bytes::format_bytes; use hg::errors::HgError; use hg::repo::Repo; use hg::utils::{files::get_bytes_from_os_str, shell_quote}; +use std::ffi::OsString; const ONE_MEBIBYTE: u64 = 1 << 20; @@ -83,14 +84,21 @@ impl<'a> Blackbox<'a> { }) } - pub fn log_command_start(&self) { + pub fn log_command_start<'arg>( + &self, + argv: impl Iterator, + ) { if let Some(configured) = &self.configured { - let message = format_bytes!(b"(rust) {}", format_cli_args()); + let message = format_bytes!(b"(rust) {}", format_cli_args(argv)); configured.log(&self.process_start_time.calendar_based, &message); } } - pub fn log_command_end(&self, exit_code: i32) { + pub fn log_command_end<'arg>( + &self, + argv: impl Iterator, + exit_code: i32, + ) { if let Some(configured) = &self.configured { let now = chrono::Local::now(); let duration = self @@ -100,7 +108,7 @@ impl<'a> Blackbox<'a> { .as_secs_f64(); let message = format_bytes!( b"(rust) {} exited {} after {} seconds", - format_cli_args(), + format_cli_args(argv), exit_code, format_bytes::Utf8(format_args!("{:.03}", duration)) ); @@ -147,8 +155,9 @@ impl ConfiguredBlackbox<'_> { } } -fn format_cli_args() -> Vec { - let mut args = std::env::args_os(); +fn format_cli_args<'a>( + mut args: impl Iterator, +) -> Vec { let _ = args.next(); // Skip the first (or zeroth) arg, the name of the `rhg` executable let mut args = args.map(|arg| shell_quote(&get_bytes_from_os_str(arg))); let mut formatted = Vec::new(); diff --git a/rust/rhg/src/main.rs b/rust/rhg/src/main.rs --- a/rust/rhg/src/main.rs +++ b/rust/rhg/src/main.rs @@ -25,6 +25,7 @@ pub mod utils { } fn main_with_result( + argv: Vec, process_start_time: &blackbox::ProcessStartTime, ui: &ui::Ui, repo: Result<&Repo, &NoRepoInCwdError>, @@ -78,7 +79,7 @@ fn main_with_result( .version("0.0.1"); let app = add_subcommand_args(app); - let matches = app.clone().get_matches_safe()?; + let matches = app.clone().get_matches_from_safe(argv.iter())?; let (subcommand_name, subcommand_matches) = matches.subcommand(); @@ -123,9 +124,9 @@ fn main_with_result( if config.is_extension_enabled(b"blackbox") { let blackbox = blackbox::Blackbox::new(&invocation, process_start_time)?; - blackbox.log_command_start(); + blackbox.log_command_start(argv.iter()); let result = run(&invocation); - blackbox.log_command_end(exit_code( + blackbox.log_command_end(argv.iter(), exit_code( &result, // TODO: show a warning or combine with original error if // `get_bool` returns an error @@ -139,7 +140,7 @@ fn main_with_result( } } -fn main() { +fn rhg_main(argv: Vec) -> ! { // Run this first, before we find out if the blackbox extension is even // enabled, in order to include everything in-between in the duration // measurements. Reading config files can be slow if they’re on NFS. @@ -147,7 +148,7 @@ fn main() { env_logger::init(); - let early_args = EarlyArgs::parse(std::env::args_os()); + let early_args = EarlyArgs::parse(&argv); let initial_current_dir = early_args.cwd.map(|cwd| { let cwd = get_path_from_bytes(&cwd); @@ -158,6 +159,7 @@ fn main() { }) .unwrap_or_else(|error| { exit( + &argv, &None, &Ui::new_infallible(&Config::empty()), OnUnsupported::Abort, @@ -179,6 +181,7 @@ fn main() { let on_unsupported = OnUnsupported::Abort; exit( + &argv, &initial_current_dir, &Ui::new_infallible(&Config::empty()), on_unsupported, @@ -191,6 +194,7 @@ fn main() { .load_cli_args(early_args.config, early_args.color) .unwrap_or_else(|error| { exit( + &argv, &initial_current_dir, &Ui::new_infallible(&non_repo_config), OnUnsupported::from_config(&non_repo_config), @@ -209,6 +213,7 @@ fn main() { } if SCHEME_RE.is_match(&repo_path_bytes) { exit( + &argv, &initial_current_dir, &Ui::new_infallible(&non_repo_config), OnUnsupported::from_config(&non_repo_config), @@ -299,6 +304,7 @@ fn main() { Err(NoRepoInCwdError { cwd: at }) } Err(error) => exit( + &argv, &initial_current_dir, &Ui::new_infallible(&non_repo_config), OnUnsupported::from_config(&non_repo_config), @@ -318,6 +324,7 @@ fn main() { }; let ui = Ui::new(&config).unwrap_or_else(|error| { exit( + &argv, &initial_current_dir, &Ui::new_infallible(&config), OnUnsupported::from_config(&config), @@ -330,12 +337,14 @@ fn main() { let on_unsupported = OnUnsupported::from_config(config); let result = main_with_result( + argv.iter().map(|s| s.to_owned()).collect(), &process_start_time, &ui, repo_result.as_ref(), config, ); exit( + &argv, &initial_current_dir, &ui, on_unsupported, @@ -348,6 +357,10 @@ fn main() { ) } +fn main() -> ! { + rhg_main(std::env::args_os().collect()) +} + fn exit_code( result: &Result<(), CommandError>, use_detailed_exit_code: bool, @@ -374,7 +387,8 @@ fn exit_code( } } -fn exit( +fn exit<'a>( + original_args: &'a [OsString], initial_current_dir: &Option, ui: &Ui, mut on_unsupported: OnUnsupported, @@ -386,7 +400,7 @@ fn exit( Err(CommandError::UnsupportedFeature { message }), ) = (&on_unsupported, &result) { - let mut args = std::env::args_os(); + let mut args = original_args.iter(); let executable = match executable { None => { exit_no_fallback( @@ -546,7 +560,7 @@ struct EarlyArgs { } impl EarlyArgs { - fn parse(args: impl IntoIterator) -> Self { + fn parse<'a>(args: impl IntoIterator) -> Self { let mut args = args.into_iter().map(get_bytes_from_os_str); let mut config = Vec::new(); let mut color = None;