From 2996220568d5dcc437b09dcfa2654a90c3ed9809 Mon Sep 17 00:00:00 2001
From: Joshua Poehls <joshua@poehls.me>
Date: Wed, 9 Sep 2020 10:52:30 -0500
Subject: [PATCH] fix(pwsh): update init script to propagate $LASTEXITCODE and
 $?

Previously the prompt function used in PowerShell would overwrite the $LASTEXITCODE and $? automatic variables that were set by the previous command run the user in the shell. This results in surprising behavior for the user if they inspect those variables looking for the result of the command they last ran.

This fixes the bug reported here: https://github.com/starship/starship/issues/1051
And goes further to also propagate the $? automatic variable which is not mentioned in that bug.
---
 src/init/starship.ps1 | 34 ++++++++++++++++++++++++++++++++--
 1 file changed, 32 insertions(+), 2 deletions(-)

diff --git a/src/init/starship.ps1 b/src/init/starship.ps1
index 87cef7c6..558d9c5c 100644
--- a/src/init/starship.ps1
+++ b/src/init/starship.ps1
@@ -3,24 +3,54 @@
 # Starship assumes UTF-8
 [Console]::OutputEncoding = [System.Text.Encoding]::UTF8
 function global:prompt {
+    $origDollarQuestion = $global:?
+    $origLastExitCode = $global:LASTEXITCODE
+
     $out = $null
     # @ makes sure the result is an array even if single or no values are returned
     $jobs = @(Get-Job | Where-Object { $_.State -eq 'Running' }).Count
 
     $env:PWD = $PWD
     $current_directory = (Convert-Path -LiteralPath $PWD)
+    
+    # If an external command has not been executed, then the $LASTEXITCODE will be $null.
+    # For the purposes of the prompt, replace this with a 0 exit code so that the prompt
+    # doesn't show the last command as a failure (because it had a non-zero exit code).
+    $lastExitCodeForPrompt = if ($origLastExitCode) { $origLastExitCode } else { 0 }
 
     if ($lastCmd = Get-History -Count 1) {
         $duration = [math]::Round(($lastCmd.EndExecutionTime - $lastCmd.StartExecutionTime).TotalMilliseconds)
         # & ensures the path is interpreted as something to execute
-        $out = @(&::STARSHIP:: prompt "--path=$current_directory" --status=$lastexitcode --jobs=$jobs --cmd-duration=$duration)
+        $out = @(&::STARSHIP:: prompt "--path=$current_directory" --status=$lastExitCodeForPrompt --jobs=$jobs --cmd-duration=$duration)
     } else {
-        $out = @(&::STARSHIP:: prompt "--path=$current_directory" --status=$lastexitcode --jobs=$jobs)
+        $out = @(&::STARSHIP:: prompt "--path=$current_directory" --status=$lastExitCodeForPrompt --jobs=$jobs)
     }
 
     # Convert stdout (array of lines) to expected return type string
     # `n is an escaped newline
     $out -join "`n"
+
+    # Propagate the original $LASTEXITCODE from before the prompt function was invoked.
+    $global:LASTEXITCODE = $origLastExitCode
+
+    # Propagate the original $? automatic variable value from before the prompt function was invoked.
+    #
+    # $? is a read-only or constant variable so we can't directly override it.
+    # In order to propagate up its original boolean value we will take an action
+    # which will produce the desired value.
+    #
+    # This has to be the very last thing that happens in the prompt function
+    # since every PowerShell command sets the $? variable.
+    if ($global:? -ne $origDollarQuestion) {
+        if ($origDollarQuestion) {
+             # Simple command which will execute successfully and set $? = True without any other side affects.
+            1+1
+        } else {
+            # Write-Error will set $? to False.
+            # ErrorAction Ignore will prevent the error from being added to the $Error collection.
+            Write-Error '' -ErrorAction 'Ignore'
+        }
+    }
 }
 
 $ENV:STARSHIP_SHELL = "powershell"