From b8bcff466bb9ed4bd7802ef74c97f89a79c284f1 Mon Sep 17 00:00:00 2001
From: prsabahrami <p.bahraminejad@gmail.com>
Date: Mon, 9 Sep 2024 12:35:47 -0400
Subject: [PATCH] Added last cmd exit_code support

---
 crates/deno_task_shell/src/parser.rs        |  4 +++-
 crates/deno_task_shell/src/shell/execute.rs |  5 +++++
 crates/deno_task_shell/src/shell/types.rs   | 10 ++++++++++
 crates/shell/src/main.rs                    |  3 ++-
 scripts/exit_code.sh                        |  1 +
 5 files changed, 21 insertions(+), 2 deletions(-)
 create mode 100644 scripts/exit_code.sh

diff --git a/crates/deno_task_shell/src/parser.rs b/crates/deno_task_shell/src/parser.rs
index 9032f53..947cb2e 100644
--- a/crates/deno_task_shell/src/parser.rs
+++ b/crates/deno_task_shell/src/parser.rs
@@ -280,6 +280,8 @@ pub enum WordPart {
   Quoted(Vec<WordPart>),
   /// Tilde prefix (ex. `~user/path` or `~/bin`)
   Tilde(TildePrefix),
+  /// Exit status of the last command (ex. `$?`)
+  ExitStatus,
 }
 
 #[cfg_attr(feature = "serialization", derive(serde::Serialize))]
@@ -704,7 +706,7 @@ fn parse_word(pair: Pair<Rule>) -> Result<Word> {
     Rule::UNQUOTED_PENDING_WORD => {
       for part in pair.into_inner() {
         match part.as_rule() {
-          Rule::EXIT_STATUS => parts.push(WordPart::Variable("?".to_string())),
+          Rule::EXIT_STATUS => parts.push(WordPart::ExitStatus),
           Rule::UNQUOTED_CHAR => {
             if let Some(WordPart::Text(ref mut text)) = parts.last_mut() {
               text.push(part.as_str().chars().next().unwrap());
diff --git a/crates/deno_task_shell/src/shell/execute.rs b/crates/deno_task_shell/src/shell/execute.rs
index 62daa5d..debbfcf 100644
--- a/crates/deno_task_shell/src/shell/execute.rs
+++ b/crates/deno_task_shell/src/shell/execute.rs
@@ -917,6 +917,11 @@ fn evaluate_word_parts(
             }
             continue;
           }
+          WordPart::ExitStatus => {
+            let exit_code = state.last_command_exit_code();
+            current_text.push(TextPart::Text(exit_code.to_string()));
+            continue;
+          }
         };
 
         // This text needs to be turned into a vector of strings.
diff --git a/crates/deno_task_shell/src/shell/types.rs b/crates/deno_task_shell/src/shell/types.rs
index ec21772..f69b977 100644
--- a/crates/deno_task_shell/src/shell/types.rs
+++ b/crates/deno_task_shell/src/shell/types.rs
@@ -36,6 +36,7 @@ pub struct ShellState {
   git_root: PathBuf, // Path to the root (`$git_root/.git/HEAD` exists)
   git_branch: String, // Contents of `$git_root/.git/HEAD`
   last_command_cd: bool, // Was last command a `cd` (thus git_branch is current)?
+  last_command_exit_code: i32, // Exit code of the last command
 }
 
 impl ShellState {
@@ -57,6 +58,7 @@ impl ShellState {
       git_root: PathBuf::new(),
       git_branch: String::new(),
       last_command_cd: false,
+      last_command_exit_code: 0,
     };
     // ensure the data is normalized
     for (name, value) in env_vars {
@@ -86,6 +88,14 @@ impl ShellState {
     self.last_command_cd
   }
 
+  pub fn set_last_command_exit_code(&mut self, exit_code: i32) {
+    self.last_command_exit_code = exit_code;
+  }
+
+  pub fn last_command_exit_code(&self) -> i32 {
+    self.last_command_exit_code
+  }
+
   pub fn env_vars(&self) -> &HashMap<String, String> {
     &self.env_vars
   }
diff --git a/crates/shell/src/main.rs b/crates/shell/src/main.rs
index 454c464..634e8f0 100644
--- a/crates/shell/src/main.rs
+++ b/crates/shell/src/main.rs
@@ -139,9 +139,10 @@ async fn interactive() -> anyhow::Result<()> {
                 rl.add_history_entry(line.as_str())?;
 
                 // Process the input (here we just echo it back)
-                _prev_exit_code = execute(&line, &mut state)
+                let prev_exit_code = execute(&line, &mut state)
                     .await
                     .context("Failed to execute")?;
+                state.set_last_command_exit_code(prev_exit_code);
 
                 // Check for exit command
                 if line.trim().eq_ignore_ascii_case("exit") {
diff --git a/scripts/exit_code.sh b/scripts/exit_code.sh
new file mode 100644
index 0000000..cc1464a
--- /dev/null
+++ b/scripts/exit_code.sh
@@ -0,0 +1 @@
+echo $?
\ No newline at end of file