Merge pull request #447 from dahlbyk/SupportsVirtualTerminal

dahlbyk · web-flow · commit d42d7a773f46 · 2017-03-03T00:14:01.000-06:00
Use ANSI color codes if SupportsVirtualTerminal
diff --git a/src/AnsiUtils.ps1 b/src/AnsiUtils.ps1
@@ -0,0 +1,36 @@
+# Color codes from https://msdn.microsoft.com/en-us/library/windows/desktop/mt638032(v=vs.85).aspx
+$ConsoleColorToAnsi = @(
+    30 # Black
+    34 # DarkBlue
+    32 # DarkGreen
+    36 # DarkCyan
+    31 # DarkRed
+    35 # DarkMagenta
+    33 # DarkYellow
+    37 # Gray
+    90 # DarkGray
+    94 # Blue
+    92 # Green
+    96 # Cyan
+    91 # Red
+    95 # Magenta
+    93 # Yellow
+    97 # White
+)
+$AnsiDefaultColor = 39
+$AnsiEscape = [char]27 + "["
+
+function Get-VirtualTerminalSequence ($color, [int]$offset = 0) {
+    if (($color -is [ConsoleColor]) -and ($color -ge 0) -and ($color -le 15)) {
+        return "${AnsiEscape}$($ConsoleColorToAnsi[$color] + $offset)m"
+    }
+    return "${AnsiEscape}$($AnsiDefaultColor + $offset)m"
+}
+
+function Get-ForegroundVirtualTerminalSequence($Color) {
+    return Get-VirtualTerminalSequence $Color
+}
+
+function Get-BackgroundVirtualTerminalSequence($Color) {
+    return Get-VirtualTerminalSequence $Color 10
+}
diff --git a/src/ConsoleMode.ps1 b/src/ConsoleMode.ps1
@@ -0,0 +1,103 @@
+# Hack! https://gist.github.com/lzybkr/f2059cb2ee8d0c13c65ab933b75e998c
+
+if ($IsWindows -eq $false) {
+    function Set-ConsoleMode {
+        [Diagnostics.CodeAnalysis.SuppressMessageAttribute("PSUseShouldProcessForStateChangingFunctions", "")]
+        param()
+    }
+    return
+}
+
+Add-Type @"
+using System;
+using System.Runtime.InteropServices;
+
+public class NativeConsoleMethods
+{
+    [DllImport("kernel32.dll", CharSet = CharSet.Unicode, SetLastError = true)]
+    public static extern IntPtr GetStdHandle(int handleId);
+
+    [DllImport("kernel32.dll", CharSet = CharSet.Unicode, SetLastError = true)]
+    public static extern bool GetConsoleMode(IntPtr hConsoleOutput, out uint dwMode);
+
+    [DllImport("kernel32.dll", CharSet = CharSet.Unicode, SetLastError = true)]
+    public static extern bool SetConsoleMode(IntPtr hConsoleOutput, uint dwMode);
+
+    public static uint GetConsoleMode(bool input = false)
+    {
+        var handle = GetStdHandle(input ? -10 : -11);
+        uint mode;
+        if (GetConsoleMode(handle, out mode))
+        {
+            return mode;
+        }
+        return 0xffffffff;
+    }
+
+    public static uint SetConsoleMode(bool input, uint mode)
+    {
+        var handle = GetStdHandle(input ? -10 : -11);
+        if (SetConsoleMode(handle, mode))
+        {
+            return GetConsoleMode(input);
+        }
+        return 0xffffffff;
+    }
+}
+"@
+
+[Flags()]
+enum ConsoleModeInputFlags
+{
+    ENABLE_PROCESSED_INPUT             = 0x0001
+    ENABLE_LINE_INPUT                  = 0x0002
+    ENABLE_ECHO_INPUT                  = 0x0004
+    ENABLE_WINDOW_INPUT                = 0x0008
+    ENABLE_MOUSE_INPUT                 = 0x0010
+    ENABLE_INSERT_MODE                 = 0x0020
+    ENABLE_QUICK_EDIT_MODE             = 0x0040
+    ENABLE_EXTENDED_FLAGS              = 0x0080
+    ENABLE_AUTO_POSITION               = 0x0100
+    ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x0200
+}
+
+[Flags()]
+enum ConsoleModeOutputFlags
+{
+    ENABLE_PROCESSED_OUTPUT            = 0x0001
+    ENABLE_WRAP_AT_EOL_OUTPUT          = 0x0002
+    ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x0004
+}
+
+function Set-ConsoleMode
+{
+    [Diagnostics.CodeAnalysis.SuppressMessageAttribute("PSUseShouldProcessForStateChangingFunctions", "")]
+    param(
+        [Parameter(ParameterSetName = "ANSI")]
+        [switch]
+        $ANSI,
+
+        [Parameter(ParameterSetName = "Mode")]
+        [uint32]
+        $Mode,
+
+        [switch]
+        $StandardInput
+    )
+
+    if ($ANSI)
+    {
+        $outputMode = [NativeConsoleMethods]::GetConsoleMode($false)
+        $null = [NativeConsoleMethods]::SetConsoleMode($false, $outputMode -bor [ConsoleModeOutputFlags]::ENABLE_VIRTUAL_TERMINAL_PROCESSING)
+
+        if ($StandardInput)
+        {
+            $inputMode = [NativeConsoleMethods]::GetConsoleMode($true)
+            $null = [NativeConsoleMethods]::SetConsoleMode($true, $inputMode -bor [ConsoleModeInputFlags]::ENABLE_VIRTUAL_TERMINAL_PROCESSING)
+        }
+    }
+    else
+    {
+        [NativeConsoleMethods]::SetConsoleMode($StandardInput, $Mode)
+    }
+}
diff --git a/src/GitPrompt.ps1 b/src/GitPrompt.ps1
@@ -97,6 +97,8 @@ $global:GitPromptSettings = [pscustomobject]@{
 
     EnableWindowTitle                           = 'posh~git ~ '
 
+    AnsiConsole                                 = $Host.UI.SupportsVirtualTerminal -or ($Env:ConEmuANSI -eq "ON")
+
     DefaultPromptPrefix                         = ''
     DefaultPromptSuffix                         = '$(''>'' * ($nestedPromptLevel + 1)) '
     DefaultPromptDebugSuffix                    = ' [DBG]$(''>'' * ($nestedPromptLevel + 1)) '
@@ -127,12 +129,35 @@ if (Get-Module NuGet) {
     $WindowTitleSupported = $false
 }
 
-function Write-Prompt($Object, $ForegroundColor, $BackgroundColor = -1) {
+function Write-Prompt {
+    param(
+        [parameter(Mandatory = $true)]
+        $Object,
+        [parameter(Mandatory = $true)]
+        $ForegroundColor,
+        [parameter()]
+        $BackgroundColor = $null,
+        [parameter(ValueFromPipeline = $true)]
+        [Text.StringBuilder]
+        $Builder)
+    if ($GitPromptSettings.AnsiConsole) {
+        $e = [char]27 + "["
+        $f = Get-ForegroundVirtualTerminalSequence $ForegroundColor
+        $b = Get-BackgroundVirtualTerminalSequence $BackgroundColor
+        if ($Builder) {
+            return $Builder.Append($f).Append($b).Append($Object).Append($e).Append("0m")
+        }
+        return "${f}${b}${Object}${e}0m"
+    }
     if ($BackgroundColor -lt 0) {
         Write-Host $Object -NoNewLine -ForegroundColor $ForegroundColor
     } else {
         Write-Host $Object -NoNewLine -ForegroundColor $ForegroundColor -BackgroundColor $BackgroundColor
     }
+    if ($Builder) {
+        return $Builder
+    }
+    return ""
 }
 
 function Format-BranchName($branchName){
@@ -148,8 +173,9 @@ function Format-BranchName($branchName){
 
 function Write-GitStatus($status) {
     $s = $global:GitPromptSettings
+    $sb = New-Object System.Text.StringBuilder
     if ($status -and $s) {
-        Write-Prompt $s.BeforeText -BackgroundColor $s.BeforeBackgroundColor -ForegroundColor $s.BeforeForegroundColor
+        $sb | Write-Prompt $s.BeforeText -BackgroundColor $s.BeforeBackgroundColor -ForegroundColor $s.BeforeForegroundColor | Out-Null
 
         $branchStatusText            = $null
         $branchStatusBackgroundColor = $s.BranchBackgroundColor
@@ -201,47 +227,47 @@ function Write-GitStatus($status) {
             $branchStatusText            = "?"
         }
 
-        Write-Prompt (Format-BranchName($status.Branch)) -BackgroundColor $branchStatusBackgroundColor -ForegroundColor $branchStatusForegroundColor
+        $sb | Write-Prompt (Format-BranchName($status.Branch)) -BackgroundColor $branchStatusBackgroundColor -ForegroundColor $branchStatusForegroundColor | Out-Null
 
         if ($branchStatusText) {
-            Write-Prompt  (" {0}" -f $branchStatusText) -BackgroundColor $branchStatusBackgroundColor -ForegroundColor $branchStatusForegroundColor
+            $sb | Write-Prompt  (" {0}" -f $branchStatusText) -BackgroundColor $branchStatusBackgroundColor -ForegroundColor $branchStatusForegroundColor | Out-Null
         }
 
         if($s.EnableFileStatus -and $status.HasIndex) {
-            Write-Prompt $s.BeforeIndexText -BackgroundColor $s.BeforeIndexBackgroundColor -ForegroundColor $s.BeforeIndexForegroundColor
+            $sb | Write-Prompt $s.BeforeIndexText -BackgroundColor $s.BeforeIndexBackgroundColor -ForegroundColor $s.BeforeIndexForegroundColor | Out-Null
 
             if($s.ShowStatusWhenZero -or $status.Index.Added) {
-                Write-Prompt (" $($s.FileAddedText)$($status.Index.Added.Count)") -BackgroundColor $s.IndexBackgroundColor -ForegroundColor $s.IndexForegroundColor
+                $sb | Write-Prompt (" $($s.FileAddedText)$($status.Index.Added.Count)") -BackgroundColor $s.IndexBackgroundColor -ForegroundColor $s.IndexForegroundColor | Out-Null
             }
             if($s.ShowStatusWhenZero -or $status.Index.Modified) {
-                Write-Prompt (" $($s.FileModifiedText)$($status.Index.Modified.Count)") -BackgroundColor $s.IndexBackgroundColor -ForegroundColor $s.IndexForegroundColor
+                $sb | Write-Prompt (" $($s.FileModifiedText)$($status.Index.Modified.Count)") -BackgroundColor $s.IndexBackgroundColor -ForegroundColor $s.IndexForegroundColor | Out-Null
             }
             if($s.ShowStatusWhenZero -or $status.Index.Deleted) {
-                Write-Prompt (" $($s.FileRemovedText)$($status.Index.Deleted.Count)") -BackgroundColor $s.IndexBackgroundColor -ForegroundColor $s.IndexForegroundColor
+                $sb | Write-Prompt (" $($s.FileRemovedText)$($status.Index.Deleted.Count)") -BackgroundColor $s.IndexBackgroundColor -ForegroundColor $s.IndexForegroundColor | Out-Null
             }
 
             if ($status.Index.Unmerged) {
-                Write-Prompt (" $($s.FileConflictedText)$($status.Index.Unmerged.Count)") -BackgroundColor $s.IndexBackgroundColor -ForegroundColor $s.IndexForegroundColor
+                $sb | Write-Prompt (" $($s.FileConflictedText)$($status.Index.Unmerged.Count)") -BackgroundColor $s.IndexBackgroundColor -ForegroundColor $s.IndexForegroundColor | Out-Null
             }
 
             if($status.HasWorking) {
-                Write-Prompt $s.DelimText -BackgroundColor $s.DelimBackgroundColor -ForegroundColor $s.DelimForegroundColor
+                $sb | Write-Prompt $s.DelimText -BackgroundColor $s.DelimBackgroundColor -ForegroundColor $s.DelimForegroundColor | Out-Null
             }
         }
 
         if($s.EnableFileStatus -and $status.HasWorking) {
             if($s.ShowStatusWhenZero -or $status.Working.Added) {
-                Write-Prompt (" $($s.FileAddedText)$($status.Working.Added.Count)") -BackgroundColor $s.WorkingBackgroundColor -ForegroundColor $s.WorkingForegroundColor
+                $sb | Write-Prompt (" $($s.FileAddedText)$($status.Working.Added.Count)") -BackgroundColor $s.WorkingBackgroundColor -ForegroundColor $s.WorkingForegroundColor | Out-Null
             }
             if($s.ShowStatusWhenZero -or $status.Working.Modified) {
-                Write-Prompt (" $($s.FileModifiedText)$($status.Working.Modified.Count)") -BackgroundColor $s.WorkingBackgroundColor -ForegroundColor $s.WorkingForegroundColor
+                $sb | Write-Prompt (" $($s.FileModifiedText)$($status.Working.Modified.Count)") -BackgroundColor $s.WorkingBackgroundColor -ForegroundColor $s.WorkingForegroundColor | Out-Null
             }
             if($s.ShowStatusWhenZero -or $status.Working.Deleted) {
-                Write-Prompt (" $($s.FileRemovedText)$($status.Working.Deleted.Count)") -BackgroundColor $s.WorkingBackgroundColor -ForegroundColor $s.WorkingForegroundColor
+                $sb | Write-Prompt (" $($s.FileRemovedText)$($status.Working.Deleted.Count)") -BackgroundColor $s.WorkingBackgroundColor -ForegroundColor $s.WorkingForegroundColor | Out-Null
             }
 
             if ($status.Working.Unmerged) {
-                Write-Prompt (" $($s.FileConflictedText)$($status.Working.Unmerged.Count)") -BackgroundColor $s.WorkingBackgroundColor -ForegroundColor $s.WorkingForegroundColor
+                $sb | Write-Prompt (" $($s.FileConflictedText)$($status.Working.Unmerged.Count)") -BackgroundColor $s.WorkingBackgroundColor -ForegroundColor $s.WorkingForegroundColor | Out-Null
             }
         }
 
@@ -263,16 +289,16 @@ function Write-GitStatus($status) {
         }
 
         if ($localStatusSymbol) {
-            Write-Prompt (" {0}" -f $localStatusSymbol) -BackgroundColor $localStatusBackgroundColor -ForegroundColor $localStatusForegroundColor
+            $sb | Write-Prompt (" {0}" -f $localStatusSymbol) -BackgroundColor $localStatusBackgroundColor -ForegroundColor $localStatusForegroundColor | Out-Null
         }
 
         if ($s.EnableStashStatus -and ($status.StashCount -gt 0)) {
-             Write-Prompt $s.BeforeStashText -BackgroundColor $s.BeforeStashBackgroundColor -ForegroundColor $s.BeforeStashForegroundColor
-             Write-Prompt $status.StashCount -BackgroundColor $s.StashBackgroundColor -ForegroundColor $s.StashForegroundColor
-             Write-Prompt $s.AfterStashText -BackgroundColor $s.AfterStashBackgroundColor -ForegroundColor $s.AfterStashForegroundColor
+             $sb | Write-Prompt $s.BeforeStashText -BackgroundColor $s.BeforeStashBackgroundColor -ForegroundColor $s.BeforeStashForegroundColor | Out-Null
+             $sb | Write-Prompt $status.StashCount -BackgroundColor $s.StashBackgroundColor -ForegroundColor $s.StashForegroundColor | Out-Null
+             $sb | Write-Prompt $s.AfterStashText -BackgroundColor $s.AfterStashBackgroundColor -ForegroundColor $s.AfterStashForegroundColor | Out-Null
         }
 
-        Write-Prompt $s.AfterText -BackgroundColor $s.AfterBackgroundColor -ForegroundColor $s.AfterForegroundColor
+        $sb | Write-Prompt $s.AfterText -BackgroundColor $s.AfterBackgroundColor -ForegroundColor $s.AfterForegroundColor | Out-Null
 
         if ($WindowTitleSupported -and $s.EnableWindowTitle) {
             if( -not $Global:PreviousWindowTitle ) {
@@ -282,8 +308,11 @@ function Write-GitStatus($status) {
             $prefix = if ($s.EnableWindowTitle -is [string]) { $s.EnableWindowTitle } else { '' }
             $Host.UI.RawUI.WindowTitle = "$script:adminHeader$prefix$repoName [$($status.Branch)]"
         }
+
+        return $sb.ToString()
     } elseif ( $Global:PreviousWindowTitle ) {
         $Host.UI.RawUI.WindowTitle = $Global:PreviousWindowTitle
+        return ""
     }
 }
 
@@ -304,6 +333,7 @@ if ($Host.UI.RawUI.BackgroundColor -eq [ConsoleColor]::DarkMagenta) {
 }
 
 function Global:Write-VcsStatus {
+    Set-ConsoleMode -ANSI
     $Global:VcsPromptStatuses | ForEach-Object { & $_ }
 }
 
diff --git a/src/posh-git.psm1 b/src/posh-git.psm1
@@ -2,7 +2,9 @@ param([switch]$NoVersionWarn, [switch]$ForcePoshGitPrompt)
 
 & $PSScriptRoot\CheckRequirements.ps1 > $null
 
+. $PSScriptRoot\ConsoleMode.ps1
 . $PSScriptRoot\Utils.ps1
+. $PSScriptRoot\AnsiUtils.ps1
 . $PSScriptRoot\GitUtils.ps1
 . $PSScriptRoot\GitPrompt.ps1
 . $PSScriptRoot\GitParamTabExpansion.ps1
@@ -70,18 +72,20 @@ if ($ForcePoshGitPrompt -or !$currentPromptDef -or ($currentPromptDef -eq $defau
             $currentPath = "~" + $currentPath.SubString($Home.Length)
         }
 
+        $res = ''
+
         # Display default prompt prefix if not empty.
         $defaultPromptPrefix = [string]$GitPromptSettings.DefaultPromptPrefix
         if ($defaultPromptPrefix) {
             $expandedDefaultPromptPrefix = $ExecutionContext.SessionState.InvokeCommand.ExpandString($defaultPromptPrefix)
-            Write-Host $expandedDefaultPromptPrefix -NoNewline
+            $res += Write-Prompt $expandedDefaultPromptPrefix -ForegroundColor $GitPromptSettings.DefaultForegroundColor
         }
 
         # Write the abbreviated current path
-        Write-Host $currentPath -NoNewline
+        $res += Write-Prompt $currentPath -ForegroundColor $GitPromptSettings.DefaultForegroundColor
 
         # Write the Git status summary information
-        Write-VcsStatus
+        $res += Write-VcsStatus
 
         # If stopped in the debugger, the prompt needs to indicate that in some fashion
         $hasInBreakpoint = [runspace]::DefaultRunspace.Debugger | Get-Member -Name InBreakpoint -MemberType property
@@ -99,11 +103,11 @@ if ($ForcePoshGitPrompt -or !$currentPromptDef -or ($currentPromptDef -eq $defau
         if ($GitPromptSettings.DefaultPromptEnableTiming) {
             $sw.Stop()
             $elapsed = $sw.ElapsedMilliseconds
-            Write-Host " ${elapsed}ms" -NoNewline
+            $res += Write-Prompt " ${elapsed}ms" -ForegroundColor $GitPromptSettings.DefaultForegroundColor
         }
 
         $global:LASTEXITCODE = $origLastExitCode
-        $expandedPromptSuffix
+        $res + $expandedPromptSuffix
     }
 
     # Set the posh-git prompt as the default prompt
diff --git a/test/DefaultPrompt.Tests.ps1 b/test/DefaultPrompt.Tests.ps1
@@ -2,6 +2,7 @@
 
 Describe 'Default Prompt Tests' {
     BeforeAll {
+        $GitPromptSettings.AnsiConsole = $false
         $prompt = Get-Item Function:\prompt
         $OFS = ''
     }

Original file line number	Diff line number	Diff line change
`@@ -2,6 +2,7 @@`
`2`	`2`
`3`	`3`	`Describe 'Default Prompt Tests' {`
`4`	`4`	`BeforeAll {`
	`5`	`+ $GitPromptSettings.AnsiConsole = $false`
`5`	`6`	`$prompt = Get-Item Function:\prompt`
`6`	`7`	`$OFS = ''`
`7`	`8`	`}`