prompt with env variable

crates/nu-cli/src/lib.rs

@@ -1,9 +1,11 @@
 mod completions;
 mod errors;
+mod prompt;
 mod syntax_highlight;
 mod validation;
 
 pub use completions::NuCompleter;
 pub use errors::report_error;
+pub use prompt::NushellPrompt;
 pub use syntax_highlight::NuHighlighter;
 pub use validation::NuValidator;

crates/nu-cli/src/prompt.rs

@@ -0,0 +1,90 @@
+use {
+    reedline::{
+        Prompt, PromptEditMode, PromptHistorySearch, PromptHistorySearchStatus, PromptViMode,
+    },
+    std::borrow::Cow,
+};
+
+/// Nushell prompt definition
+#[derive(Clone)]
+pub struct NushellPrompt {
+    prompt_command: String,
+    prompt_string: String,
+    // These are part of the struct definition in case we want to later allow
+    // further customization to the shell status
+    default_prompt_indicator: String,
+    default_vi_insert_prompt_indicator: String,
+    default_vi_visual_prompt_indicator: String,
+    default_multiline_indicator: String,
+}
+
+impl Default for NushellPrompt {
+    fn default() -> Self {
+        NushellPrompt::new()
+    }
+}
+
+impl NushellPrompt {
+    /// Constructor for the default prompt, which takes the amount of spaces required between the left and right-hand sides of the prompt
+    pub fn new() -> NushellPrompt {
+        NushellPrompt {
+            prompt_command: "".to_string(),
+            prompt_string: "".to_string(),
+            default_prompt_indicator: "〉".to_string(),
+            default_vi_insert_prompt_indicator: ": ".to_string(),
+            default_vi_visual_prompt_indicator: "v ".to_string(),
+            default_multiline_indicator: "::: ".to_string(),
+        }
+    }
+
+    pub fn is_new_prompt(&self, prompt_command: &str) -> bool {
+        self.prompt_command != prompt_command
+    }
+
+    pub fn update_prompt(&mut self, prompt_command: String, prompt_string: String) {
+        self.prompt_command = prompt_command;
+        self.prompt_string = prompt_string;
+    }
+
+    fn default_wrapped_custom_string(&self, str: String) -> String {
+        format!("({})", str)
+    }
+}
+
+impl Prompt for NushellPrompt {
+    fn render_prompt(&self, _: usize) -> Cow<str> {
+        self.prompt_string.as_str().into()
+    }
+
+    fn render_prompt_indicator(&self, edit_mode: PromptEditMode) -> Cow<str> {
+        match edit_mode {
+            PromptEditMode::Default => self.default_prompt_indicator.as_str().into(),
+            PromptEditMode::Emacs => self.default_prompt_indicator.as_str().into(),
+            PromptEditMode::Vi(vi_mode) => match vi_mode {
+                PromptViMode::Normal => self.default_prompt_indicator.as_str().into(),
+                PromptViMode::Insert => self.default_vi_insert_prompt_indicator.as_str().into(),
+                PromptViMode::Visual => self.default_vi_visual_prompt_indicator.as_str().into(),
+            },
+            PromptEditMode::Custom(str) => self.default_wrapped_custom_string(str).into(),
+        }
+    }
+
+    fn render_prompt_multiline_indicator(&self) -> Cow<str> {
+        Cow::Borrowed(self.default_multiline_indicator.as_str())
+    }
+
+    fn render_prompt_history_search_indicator(
+        &self,
+        history_search: PromptHistorySearch,
+    ) -> Cow<str> {
+        let prefix = match history_search.status {
+            PromptHistorySearchStatus::Passing => "",
+            PromptHistorySearchStatus::Failing => "failing ",
+        };
+
+        Cow::Owned(format!(
+            "({}reverse-search: {})",
+            prefix, history_search.term
+        ))
+    }
+}

crates/nu-protocol/src/engine/evaluation_context.rs

@@ -112,6 +112,10 @@ impl Stack {
         self.0.borrow().env_vars.clone()
     }
 
+    pub fn get_env_var(&self, name: &str) -> Option<String> {
+        self.0.borrow().env_vars.get(name).cloned()
+    }
+
     pub fn print_stack(&self) {
         println!("===frame===");
         println!("vars:");

src/main.rs

@@ -1,20 +1,23 @@
-use std::io::Write;
+use std::{cell::RefCell, io::Write, rc::Rc};
 
 use miette::{IntoDiagnostic, Result};
-use nu_cli::{report_error, NuCompleter, NuHighlighter, NuValidator};
+use nu_cli::{report_error, NuCompleter, NuHighlighter, NuValidator, NushellPrompt};
 use nu_command::create_default_context;
 use nu_engine::eval_block;
 use nu_parser::parse;
 use nu_protocol::{
     ast::Call,
-    engine::{EngineState, EvaluationContext, StateWorkingSet},
+    engine::{EngineState, EvaluationContext, Stack, StateWorkingSet},
     ShellError, Value,
 };
-use reedline::DefaultCompletionActionHandler;
+use reedline::{DefaultPrompt, Prompt};
 
 #[cfg(test)]
 mod tests;
 
+// Name of environment variable where the prompt could be stored
+const PROMPT_COMMAND: &str = "PROMPT_COMMAND";
+
 fn main() -> Result<()> {
     miette::set_panic_hook();
     let miette_hook = std::panic::take_hook();
@@ -63,7 +66,7 @@ fn main() -> Result<()> {
 
         Ok(())
     } else {
-        use reedline::{DefaultPrompt, FileBackedHistory, Reedline, Signal};
+        use reedline::{DefaultCompletionActionHandler, FileBackedHistory, Reedline, Signal};
 
         let completer = NuCompleter::new(engine_state.clone());
         let mut entry_num = 0;
@@ -84,13 +87,22 @@ fn main() -> Result<()> {
                 engine_state: engine_state.clone(),
             }));
 
-        let prompt = DefaultPrompt::new(1);
+        let default_prompt = DefaultPrompt::new(1);
+        let mut nu_prompt = NushellPrompt::new();
         let stack = nu_protocol::engine::Stack::new();
 
         loop {
+            let prompt = update_prompt(
+                PROMPT_COMMAND,
+                engine_state.clone(),
+                &stack,
+                &mut nu_prompt,
+                &default_prompt,
+            );
+
             entry_num += 1;
 
-            let input = line_editor.read_line(&prompt);
+            let input = line_editor.read_line(prompt);
             match input {
                 Ok(Signal::Success(s)) => {
                     if s.trim() == "exit" {
@@ -189,3 +201,62 @@ fn print_value(value: Value, state: &EvaluationContext) -> Result<(), ShellError
 
     Ok(())
 }
+
+fn update_prompt<'prompt>(
+    env_variable: &str,
+    engine_state: Rc<RefCell<EngineState>>,
+    stack: &Stack,
+    nu_prompt: &'prompt mut NushellPrompt,
+    default_prompt: &'prompt DefaultPrompt,
+) -> &'prompt dyn Prompt {
+    let prompt_command = match stack.get_env_var(env_variable) {
+        Some(prompt) => prompt,
+        None => return default_prompt as &dyn Prompt,
+    };
+
+    // Checking if the PROMPT_COMMAND is the same to avoid evaluating constantly
+    // the same command, thus saturating the contents in the EngineState
+    if !nu_prompt.is_new_prompt(prompt_command.as_str()) {
+        return nu_prompt as &dyn Prompt;
+    }
+
+    let (block, delta) = {
+        let ref_engine_state = engine_state.borrow();
+        let mut working_set = StateWorkingSet::new(&ref_engine_state);
+        let (output, err) = parse(&mut working_set, None, prompt_command.as_bytes(), false);
+        if let Some(err) = err {
+            report_error(&working_set, &err);
+            return default_prompt as &dyn Prompt;
+        }
+        (output, working_set.render())
+    };
+
+    EngineState::merge_delta(&mut *engine_state.borrow_mut(), delta);
+
+    let state = nu_protocol::engine::EvaluationContext {
+        engine_state: engine_state.clone(),
+        stack: stack.clone(),
+    };
+
+    let evaluated_prompt = match eval_block(&state, &block, Value::nothing()) {
+        Ok(value) => match value.as_string() {
+            Ok(prompt) => prompt,
+            Err(err) => {
+                let engine_state = engine_state.borrow();
+                let working_set = StateWorkingSet::new(&*engine_state);
+                report_error(&working_set, &err);
+                return default_prompt as &dyn Prompt;
+            }
+        },
+        Err(err) => {
+            let engine_state = engine_state.borrow();
+            let working_set = StateWorkingSet::new(&*engine_state);
+            report_error(&working_set, &err);
+            return default_prompt as &dyn Prompt;
+        }
+    };
+
+    nu_prompt.update_prompt(prompt_command, evaluated_prompt);
+
+    nu_prompt as &dyn Prompt
+}