From d16cca23acf2bcae05cecdb2f6bf85fe3b7cb1cb Mon Sep 17 00:00:00 2001 From: sigoden Date: Sat, 18 May 2024 22:34:23 +0800 Subject: [PATCH] refactor: shell detect/prompt (#519) --- src/config/role.rs | 15 ++++---------- src/main.rs | 25 ++++++++++------------- src/utils/command.rs | 48 ++++++++++++++++++++++++++------------------ 3 files changed, 44 insertions(+), 44 deletions(-) diff --git a/src/config/role.rs b/src/config/role.rs index 27858a3..c477ee4 100644 --- a/src/config/role.rs +++ b/src/config/role.rs @@ -225,23 +225,16 @@ fn parse_structure_prompt(prompt: &str) -> (&str, Vec<(&str, &str)>) { fn shell_prompt() -> String { let os = detect_os(); - let (detected_shell, _, _) = detect_shell(); - let (shell, use_semicolon) = match (detected_shell.as_str(), os.as_str()) { - // GPT doesn’t know much about nushell - ("nushell", "windows") => ("cmd", true), - ("nushell", _) => ("bash", true), - ("powershell", _) => ("powershell", true), - ("pwsh", _) => ("powershell", false), - _ => (detected_shell.as_str(), false), - }; - let combine = if use_semicolon { + let shell = detect_shell(); + let shell = shell.name.as_str(); + let combinator = if shell == "powershell" { "\nIf multiple steps required try to combine them together using ';'.\nIf it already combined with '&&' try to replace it with ';'.".to_string() } else { "\nIf multiple steps required try to combine them together using '&&'.".to_string() }; format!( r#"Provide only {shell} commands for {os} without any description. -Ensure the output is a valid {shell} command. {combine} +Ensure the output is a valid {shell} command. {combinator} If there is a lack of details, provide most logical solution. Output plain text only, without any markdown formatting."# ) diff --git a/src/main.rs b/src/main.rs index ca4b298..947ab37 100644 --- a/src/main.rs +++ b/src/main.rs @@ -21,7 +21,10 @@ use crate::config::{ use crate::function::eval_tool_calls; use crate::render::{render_error, MarkdownRender}; use crate::repl::Repl; -use crate::utils::{create_abort_signal, extract_block, run_command, run_spinner, CODE_BLOCK_RE}; +use crate::utils::{ + create_abort_signal, detect_shell, extract_block, run_command, run_spinner, Shell, + CODE_BLOCK_RE, +}; use anyhow::{bail, Result}; use async_recursion::async_recursion; @@ -34,7 +37,6 @@ use std::io::{stderr, stdin, stdout, Read}; use std::process; use std::sync::Arc; use tokio::sync::oneshot; -use utils::detect_shell; #[tokio::main] async fn main() -> Result<()> { @@ -118,8 +120,8 @@ async fn main() -> Result<()> { bail!("No input"); } let input = create_input(&config, text, file)?; - let (_, shell, shell_arg) = detect_shell(); - shell_execute(&config, &shell, shell_arg, input).await?; + let shell = detect_shell(); + shell_execute(&config, &shell, input).await?; return Ok(()); } config.write().apply_prelude()?; @@ -195,12 +197,7 @@ async fn start_interactive(config: &GlobalConfig) -> Result<()> { } #[async_recursion::async_recursion] -async fn shell_execute( - config: &GlobalConfig, - shell: &str, - shell_arg: &str, - mut input: Input, -) -> Result<()> { +async fn shell_execute(config: &GlobalConfig, shell: &Shell, mut input: Input) -> Result<()> { let client = input.create_client()?; let is_terminal_stdout = stdout().is_terminal(); let ret = if is_terminal_stdout { @@ -227,15 +224,15 @@ async fn shell_execute( if is_terminal_stdout { loop { let answer = Select::new( - markdown_render.render(&eval_str).trim(), + eval_str.trim(), vec!["βœ… Execute", "πŸ€” Revise", "πŸ“™ Explain", "❌ Cancel"], ) .prompt()?; match answer { "βœ… Execute" => { - debug!("{} {:?}", shell, &[shell_arg, &eval_str]); - let code = run_command(shell, &[shell_arg, &eval_str], None)?; + debug!("{} {:?}", shell.cmd, &[&shell.arg, &eval_str]); + let code = run_command(&shell.cmd, &[&shell.arg, &eval_str], None)?; if code != 0 { process::exit(code); } @@ -244,7 +241,7 @@ async fn shell_execute( let revision = Text::new("Enter your revision:").prompt()?; let text = format!("{}\n{revision}", input.text()); input.set_text(text); - return shell_execute(config, shell, shell_arg, input).await; + return shell_execute(config, shell, input).await; } "πŸ“™ Explain" => { let role = config.read().retrieve_role(EXPLAIN_SHELL_ROLE)?; diff --git a/src/utils/command.rs b/src/utils/command.rs index 77ca813..873e40b 100644 --- a/src/utils/command.rs +++ b/src/utils/command.rs @@ -16,18 +16,32 @@ pub fn detect_os() -> String { os.to_string() } -pub fn detect_shell() -> (String, String, &'static str) { +pub struct Shell { + pub name: String, + pub cmd: String, + pub arg: String, +} + +impl Shell { + pub fn new(name: &str, cmd: &str, arg: &str) -> Self { + Self { + name: name.to_string(), + cmd: cmd.to_string(), + arg: arg.to_string(), + } + } +} + +pub fn detect_shell() -> Shell { let os = env::consts::OS; if os == "windows" { - if env::var("NU_VERSION").is_ok() { - ("nushell".into(), "nu.exe".into(), "-c") - } else if let Some(ret) = env::var("PSModulePath").ok().and_then(|v| { + if let Some(ret) = env::var("PSModulePath").ok().and_then(|v| { let v = v.to_lowercase(); if v.split(';').count() >= 3 { if v.contains("powershell\\7\\") { - Some(("pwsh".into(), "pwsh.exe".into(), "-c")) + Some(Shell::new("pwsh", "pwsh.exe", "-c")) } else { - Some(("powershell".into(), "powershell.exe".into(), "-Command")) + Some(Shell::new("powershell", "powershell.exe", "-Command")) } } else { None @@ -35,22 +49,18 @@ pub fn detect_shell() -> (String, String, &'static str) { }) { ret } else { - ("cmd".into(), "cmd.exe".into(), "/C") + Shell::new("cmd", "cmd.exe", "/C") } - } else if env::var("NU_VERSION").is_ok() { - ("nushell".into(), "nu".into(), "-c") } else { - let shell_cmd = env::var("SHELL").unwrap_or_else(|_| "/bin/sh".to_string()); - let shell_name = match shell_cmd.rsplit_once('/') { - Some((_, name)) => name.to_string(), - None => shell_cmd.clone(), - }; - let shell_name = if shell_name == "nu" { - "nushell".into() - } else { - shell_name + let shell = env::var("SHELL").unwrap_or_else(|_| "/bin/sh".to_string()); + let shell = match shell.rsplit_once('/') { + Some((_, v)) => v, + None => &shell, }; - (shell_name, shell_cmd, "-c") + match shell { + "bash" | "zsh" | "fish" | "pwsh" => Shell::new(shell, shell, "-c"), + _ => Shell::new("sh", "sh", "-c"), + } } }