From 24aa1f312aa460a7db7bb04672b5173a9614b98f Mon Sep 17 00:00:00 2001 From: Hofer-Julian <30049909+Hofer-Julian@users.noreply.github.com> Date: Fri, 20 Jan 2023 19:44:49 +0100 Subject: [PATCH] Cleanup of `src/main.rs` (#7801) While reading nushell's code, I've started to clean up `src/main.rs` a bit. The description of the changes can be found in the commit messages. --- src/main.rs | 444 +++++++++++++++++++++++-------------------------- src/signals.rs | 29 ++++ 2 files changed, 236 insertions(+), 237 deletions(-) create mode 100644 src/signals.rs diff --git a/src/main.rs b/src/main.rs index fb9898cb2b..ebfe906a80 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,6 @@ mod config_files; mod logger; +mod signals; mod test_bins; #[cfg(test)] mod tests; @@ -27,12 +28,10 @@ use nu_protocol::{ SyntaxShape, Value, }; use nu_utils::stdout_write_all_and_flush; +use signals::{ctrlc_protection, sigquit_protection}; use std::{ io::BufReader, - sync::{ - atomic::{AtomicBool, Ordering}, - Arc, - }, + sync::{atomic::AtomicBool, Arc}, }; use std::{path::Path, str::FromStr}; @@ -136,7 +135,6 @@ fn acquire_terminal(interactive: bool) { fn acquire_terminal(_: bool) {} fn main() -> Result<()> { - // miette::set_panic_hook(); let miette_hook = std::panic::take_hook(); std::panic::set_hook(Box::new(move |x| { crossterm::terminal::disable_raw_mode().expect("unable to disable raw mode"); @@ -152,7 +150,6 @@ fn main() -> Result<()> { let mut working_set = nu_protocol::engine::StateWorkingSet::new(&engine_state); working_set.add_decl(Box::new(nu_cli::NuHighlight)); working_set.add_decl(Box::new(nu_cli::Print)); - working_set.render() }; @@ -160,29 +157,10 @@ fn main() -> Result<()> { report_error_new(&engine_state, &err); } - // TODO: make this conditional in the future - // Ctrl-c protection section let ctrlc = Arc::new(AtomicBool::new(false)); - let handler_ctrlc = ctrlc.clone(); - let engine_state_ctrlc = ctrlc.clone(); - - ctrlc::set_handler(move || { - handler_ctrlc.store(true, Ordering::SeqCst); - }) - .expect("Error setting Ctrl-C handler"); - - engine_state.ctrlc = Some(engine_state_ctrlc); - // End ctrl-c protection section - - // SIGQUIT protection section (only works for POSIX system) - #[cfg(not(windows))] - { - use signal_hook::consts::SIGQUIT; - let sig_quit = Arc::new(AtomicBool::new(false)); - signal_hook::flag::register(SIGQUIT, sig_quit.clone()).expect("Error setting SIGQUIT flag"); - engine_state.set_sig_quit(sig_quit); - } - // End SIGQUIT protection section + // TODO: make this conditional in the future + ctrlc_protection(&mut engine_state, &ctrlc); + sigquit_protection(&mut engine_state); let mut args_to_nushell = vec![]; let mut script_name = String::new(); @@ -231,226 +209,218 @@ fn main() -> Result<()> { let nushell_commandline_args = args_to_nushell.join(" "); - let parsed_nu_cli_args = parse_commandline_args(&nushell_commandline_args, &mut engine_state); + let parsed_nu_cli_args = parse_commandline_args(&nushell_commandline_args, &mut engine_state) + .unwrap_or_else(|_| std::process::exit(1)); - if let Ok(ref args) = parsed_nu_cli_args { - set_config_path( - &mut engine_state, - &init_cwd, - "config.nu", - "config-path", - &args.config_file, - ); + set_config_path( + &mut engine_state, + &init_cwd, + "config.nu", + "config-path", + &parsed_nu_cli_args.config_file, + ); - set_config_path( - &mut engine_state, - &init_cwd, - "env.nu", - "env-path", - &args.env_file, - ); + set_config_path( + &mut engine_state, + &init_cwd, + "env.nu", + "env-path", + &parsed_nu_cli_args.env_file, + ); + + // keep this condition in sync with the branches below + acquire_terminal( + parsed_nu_cli_args.commands.is_none() + && (script_name.is_empty() || parsed_nu_cli_args.interactive_shell.is_some()), + ); + + if let Some(t) = parsed_nu_cli_args.threads { + // 0 means to let rayon decide how many threads to use + let threads = t.as_i64().unwrap_or(0); + rayon::ThreadPoolBuilder::new() + .num_threads(threads as usize) + .build_global() + .expect("error setting number of threads"); } - match parsed_nu_cli_args { - Ok(binary_args) => { - // keep this condition in sync with the branches below - acquire_terminal( - binary_args.commands.is_none() - && (script_name.is_empty() || binary_args.interactive_shell.is_some()), + if let Some(level) = parsed_nu_cli_args.log_level.map(|level| level.item) { + let level = if Level::from_str(&level).is_ok() { + level + } else { + eprintln!( + "ERROR: log library did not recognize log level '{level}', using default 'info'" ); + "info".to_string() + }; + let target = parsed_nu_cli_args + .log_target + .map(|target| target.item) + .unwrap_or_else(|| "stderr".to_string()); - if let Some(t) = binary_args.threads { - // 0 means to let rayon decide how many threads to use - let threads = t.as_i64().unwrap_or(0); - rayon::ThreadPoolBuilder::new() - .num_threads(threads as usize) - .build_global() - .expect("error setting number of threads"); - } + logger(|builder| configure(&level, &target, builder))?; + info!("start logging {}:{}:{}", file!(), line!(), column!()); + } - if binary_args.log_level.is_some() { - let mut level = binary_args - .log_level - .map(|level| level.item) - .unwrap_or_else(|| "info".to_string()); + if let Some(testbin) = &parsed_nu_cli_args.testbin { + // Call out to the correct testbin + match testbin.item.as_str() { + "echo_env" => test_bins::echo_env(true), + "echo_env_stderr" => test_bins::echo_env(false), + "cococo" => test_bins::cococo(), + "meow" => test_bins::meow(), + "meowb" => test_bins::meowb(), + "relay" => test_bins::relay(), + "iecho" => test_bins::iecho(), + "fail" => test_bins::fail(), + "nonu" => test_bins::nonu(), + "chop" => test_bins::chop(), + "repeater" => test_bins::repeater(), + "nu_repl" => test_bins::nu_repl(), + _ => std::process::exit(1), + } + std::process::exit(0) + } + let input = if let Some(redirect_stdin) = &parsed_nu_cli_args.redirect_stdin { + let stdin = std::io::stdin(); + let buf_reader = BufReader::new(stdin); - if Level::from_str(level.as_str()).is_err() { - eprintln!("ERROR: log library did not recognize log level '{level}', using default 'info'"); - level = "info".to_string(); + PipelineData::ExternalStream { + stdout: Some(RawStream::new( + Box::new(BufferedReader::new(buf_reader)), + Some(ctrlc), + redirect_stdin.span, + None, + )), + stderr: None, + exit_code: None, + span: redirect_stdin.span, + metadata: None, + trim_end_newline: false, + } + } else { + PipelineData::empty() + }; + + info!("redirect_stdin {}:{}:{}", file!(), line!(), column!()); + + // First, set up env vars as strings only + gather_parent_env_vars(&mut engine_state, &init_cwd); + + let mut stack = nu_protocol::engine::Stack::new(); + + if let Some(commands) = &parsed_nu_cli_args.commands { + #[cfg(feature = "plugin")] + read_plugin_file( + &mut engine_state, + &mut stack, + parsed_nu_cli_args.plugin_file, + NUSHELL_FOLDER, + ); + + // only want to load config and env if relative argument is provided. + if parsed_nu_cli_args.env_file.is_some() { + config_files::read_config_file( + &mut engine_state, + &mut stack, + parsed_nu_cli_args.env_file, + true, + ); + } else { + config_files::read_default_env_file(&mut engine_state, &mut stack) + } + + if parsed_nu_cli_args.config_file.is_some() { + config_files::read_config_file( + &mut engine_state, + &mut stack, + parsed_nu_cli_args.config_file, + false, + ); + } + + let ret_val = evaluate_commands( + commands, + &mut engine_state, + &mut stack, + input, + parsed_nu_cli_args.table_mode, + ); + info!("-c command execution {}:{}:{}", file!(), line!(), column!()); + match ret_val { + Ok(Some(exit_code)) => std::process::exit(exit_code as i32), + Ok(None) => Ok(()), + Err(e) => Err(e), + } + } else if !script_name.is_empty() && parsed_nu_cli_args.interactive_shell.is_none() { + #[cfg(feature = "plugin")] + read_plugin_file( + &mut engine_state, + &mut stack, + parsed_nu_cli_args.plugin_file, + NUSHELL_FOLDER, + ); + + // only want to load config and env if relative argument is provided. + if parsed_nu_cli_args.env_file.is_some() { + config_files::read_config_file( + &mut engine_state, + &mut stack, + parsed_nu_cli_args.env_file, + true, + ); + } else { + config_files::read_default_env_file(&mut engine_state, &mut stack) + } + + if parsed_nu_cli_args.config_file.is_some() { + config_files::read_config_file( + &mut engine_state, + &mut stack, + parsed_nu_cli_args.config_file, + false, + ); + } + + let ret_val = evaluate_file( + script_name, + &args_to_script, + &mut engine_state, + &mut stack, + input, + ); + + let last_exit_code = stack.get_env_var(&engine_state, "LAST_EXIT_CODE"); + if let Some(last_exit_code) = last_exit_code { + let value = last_exit_code.as_integer(); + if let Ok(value) = value { + if value != 0 { + std::process::exit(value as i32); } - - let target = binary_args - .log_target - .map(|target| target.item) - .unwrap_or_else(|| "stderr".to_string()); - - logger(|builder| configure(level.as_str(), target.as_str(), builder))?; - info!("start logging {}:{}:{}", file!(), line!(), column!()); - } - - if let Some(testbin) = &binary_args.testbin { - // Call out to the correct testbin - match testbin.item.as_str() { - "echo_env" => test_bins::echo_env(true), - "echo_env_stderr" => test_bins::echo_env(false), - "cococo" => test_bins::cococo(), - "meow" => test_bins::meow(), - "meowb" => test_bins::meowb(), - "relay" => test_bins::relay(), - "iecho" => test_bins::iecho(), - "fail" => test_bins::fail(), - "nonu" => test_bins::nonu(), - "chop" => test_bins::chop(), - "repeater" => test_bins::repeater(), - "nu_repl" => test_bins::nu_repl(), - _ => std::process::exit(1), - } - std::process::exit(0) - } - let input = if let Some(redirect_stdin) = &binary_args.redirect_stdin { - let stdin = std::io::stdin(); - let buf_reader = BufReader::new(stdin); - - PipelineData::ExternalStream { - stdout: Some(RawStream::new( - Box::new(BufferedReader::new(buf_reader)), - Some(ctrlc), - redirect_stdin.span, - None, - )), - stderr: None, - exit_code: None, - span: redirect_stdin.span, - metadata: None, - trim_end_newline: false, - } - } else { - PipelineData::empty() - }; - - info!("redirect_stdin {}:{}:{}", file!(), line!(), column!()); - - // First, set up env vars as strings only - gather_parent_env_vars(&mut engine_state, &init_cwd); - - let mut stack = nu_protocol::engine::Stack::new(); - - if let Some(commands) = &binary_args.commands { - #[cfg(feature = "plugin")] - read_plugin_file( - &mut engine_state, - &mut stack, - binary_args.plugin_file, - NUSHELL_FOLDER, - ); - - // only want to load config and env if relative argument is provided. - if binary_args.env_file.is_some() { - config_files::read_config_file( - &mut engine_state, - &mut stack, - binary_args.env_file, - true, - ); - } else { - config_files::read_default_env_file(&mut engine_state, &mut stack) - } - - if binary_args.config_file.is_some() { - config_files::read_config_file( - &mut engine_state, - &mut stack, - binary_args.config_file, - false, - ); - } - - let ret_val = evaluate_commands( - commands, - &mut engine_state, - &mut stack, - input, - binary_args.table_mode, - ); - info!("-c command execution {}:{}:{}", file!(), line!(), column!()); - match ret_val { - Ok(Some(exit_code)) => std::process::exit(exit_code as i32), - Ok(None) => Ok(()), - Err(e) => Err(e), - } - } else if !script_name.is_empty() && binary_args.interactive_shell.is_none() { - #[cfg(feature = "plugin")] - read_plugin_file( - &mut engine_state, - &mut stack, - binary_args.plugin_file, - NUSHELL_FOLDER, - ); - - // only want to load config and env if relative argument is provided. - if binary_args.env_file.is_some() { - config_files::read_config_file( - &mut engine_state, - &mut stack, - binary_args.env_file, - true, - ); - } else { - config_files::read_default_env_file(&mut engine_state, &mut stack) - } - - if binary_args.config_file.is_some() { - config_files::read_config_file( - &mut engine_state, - &mut stack, - binary_args.config_file, - false, - ); - } - - let ret_val = evaluate_file( - script_name, - &args_to_script, - &mut engine_state, - &mut stack, - input, - ); - - let last_exit_code = stack.get_env_var(&engine_state, "LAST_EXIT_CODE"); - if let Some(last_exit_code) = last_exit_code { - let value = last_exit_code.as_integer(); - if let Ok(value) = value { - if value != 0 { - std::process::exit(value as i32); - } - } - } - info!("eval_file execution {}:{}:{}", file!(), line!(), column!()); - - ret_val - } else { - setup_config( - &mut engine_state, - &mut stack, - #[cfg(feature = "plugin")] - binary_args.plugin_file, - binary_args.config_file, - binary_args.env_file, - binary_args.login_shell.is_some(), - ); - - let ret_val = evaluate_repl( - &mut engine_state, - &mut stack, - config_files::NUSHELL_FOLDER, - binary_args.execute, - ); - info!("repl eval {}:{}:{}", file!(), line!(), column!()); - - ret_val } } - Err(_) => std::process::exit(1), + info!("eval_file execution {}:{}:{}", file!(), line!(), column!()); + + ret_val + } else { + setup_config( + &mut engine_state, + &mut stack, + #[cfg(feature = "plugin")] + parsed_nu_cli_args.plugin_file, + parsed_nu_cli_args.config_file, + parsed_nu_cli_args.env_file, + parsed_nu_cli_args.login_shell.is_some(), + ); + + let ret_val = evaluate_repl( + &mut engine_state, + &mut stack, + config_files::NUSHELL_FOLDER, + parsed_nu_cli_args.execute, + ); + info!("repl eval {}:{}:{}", file!(), line!(), column!()); + + ret_val } } diff --git a/src/signals.rs b/src/signals.rs new file mode 100644 index 0000000000..1e71789824 --- /dev/null +++ b/src/signals.rs @@ -0,0 +1,29 @@ +use std::sync::{ + atomic::{AtomicBool, Ordering}, + Arc, +}; + +use nu_protocol::engine::EngineState; + +pub(crate) fn ctrlc_protection(engine_state: &mut EngineState, ctrlc: &Arc) { + let handler_ctrlc = ctrlc.clone(); + let engine_state_ctrlc = ctrlc.clone(); + + ctrlc::set_handler(move || { + handler_ctrlc.store(true, Ordering::SeqCst); + }) + .expect("Error setting Ctrl-C handler"); + + engine_state.ctrlc = Some(engine_state_ctrlc); +} + +#[cfg(not(windows))] +pub(crate) fn sigquit_protection(engine_state: &mut EngineState) { + use signal_hook::consts::SIGQUIT; + let sig_quit = Arc::new(AtomicBool::new(false)); + signal_hook::flag::register(SIGQUIT, sig_quit.clone()).expect("Error setting SIGQUIT flag"); + engine_state.set_sig_quit(sig_quit); +} + +#[cfg(windows)] +pub(crate) fn sigquit_protection(_engine_state: &mut EngineState) {}