|
| 1 | +# Atuin PowerShell module |
| 2 | +# |
| 3 | +# Usage: atuin init powershell | Out-String | Invoke-Expression |
| 4 | + |
| 5 | +if (Get-Module Atuin -ErrorAction Ignore) { |
| 6 | + Write-Warning "The Atuin module is already loaded." |
| 7 | + return |
| 8 | +} |
| 9 | + |
| 10 | +if (!(Get-Command atuin -ErrorAction Ignore)) { |
| 11 | + Write-Error "The 'atuin' executable needs to be available in the PATH." |
| 12 | + return |
| 13 | +} |
| 14 | + |
| 15 | +if (!(Get-Module PSReadLine -ErrorAction Ignore)) { |
| 16 | + Write-Error "Atuin requires the PSReadLine module to be installed." |
| 17 | + return |
| 18 | +} |
| 19 | + |
| 20 | +New-Module -Name Atuin -ScriptBlock { |
| 21 | + $env:ATUIN_SESSION = atuin uuid |
| 22 | + |
| 23 | + $script:atuinHistoryId = $null |
| 24 | + $script:previousPSConsoleHostReadLine = $Function:PSConsoleHostReadLine |
| 25 | + |
| 26 | + # The ReadLine overloads changed with breaking changes over time, make sure the one we expect is available. |
| 27 | + $script:hasExpectedReadLineOverload = ([Microsoft.PowerShell.PSConsoleReadLine]::ReadLine).OverloadDefinitions.Contains("static string ReadLine(runspace runspace, System.Management.Automation.EngineIntrinsics engineIntrinsics, System.Threading.CancellationToken cancellationToken, System.Nullable[bool] lastRunStatus)") |
| 28 | + |
| 29 | + function PSConsoleHostReadLine { |
| 30 | + # This needs to be done as the first thing because any script run will flush $?. |
| 31 | + $lastRunStatus = $? |
| 32 | + |
| 33 | + # Exit statuses are maintained separately for native and PowerShell commands, this needs to be taken into account. |
| 34 | + $exitCode = if ($lastRunStatus) { 0 } elseif ($global:LASTEXITCODE) { $global:LASTEXITCODE } else { 1 } |
| 35 | + |
| 36 | + if ($script:atuinHistoryId) { |
| 37 | + # The duration is not recorded in old PowerShell versions, let Atuin handle it. |
| 38 | + $duration = (Get-History -Count 1).Duration.Ticks * 100 |
| 39 | + $durationArg = if ($duration) { "--duration=$duration" } else { "" } |
| 40 | + |
| 41 | + atuin history end --exit=$exitCode $durationArg -- $script:atuinHistoryId | Out-Null |
| 42 | + |
| 43 | + $global:LASTEXITCODE = $exitCode |
| 44 | + $script:atuinHistoryId = $null |
| 45 | + } |
| 46 | + |
| 47 | + # PSConsoleHostReadLine implementation from PSReadLine, adjusted to support old versions. |
| 48 | + Microsoft.PowerShell.Core\Set-StrictMode -Off |
| 49 | + |
| 50 | + $line = if ($script:hasExpectedReadLineOverload) { |
| 51 | + # When the overload we expect is available, we can pass $lastRunStatus to it. |
| 52 | + [Microsoft.PowerShell.PSConsoleReadLine]::ReadLine($Host.Runspace, $ExecutionContext, [System.Threading.CancellationToken]::None, $lastRunStatus) |
| 53 | + } else { |
| 54 | + # Either PSReadLine is older than v2.2.0-beta3, or maybe newer than we expect, so use the function from PSReadLine as-is. |
| 55 | + & $script:previousPSConsoleHostReadLine |
| 56 | + } |
| 57 | + |
| 58 | + $script:atuinHistoryId = atuin history start -- $line |
| 59 | + |
| 60 | + return $line |
| 61 | + } |
| 62 | + |
| 63 | + function RunSearch { |
| 64 | + param([string]$ExtraArgs = "") |
| 65 | + |
| 66 | + $line = $null |
| 67 | + [Microsoft.PowerShell.PSConsoleReadLine]::GetBufferState([ref]$line, [ref]$null) |
| 68 | + |
| 69 | + # Atuin is started through Start-Process to avoid interfering with the current shell, |
| 70 | + # and to capture its output which is provided in stderr (redirected to a temporary file). |
| 71 | + |
| 72 | + $suggestion = "" |
| 73 | + $resultFile = New-TemporaryFile |
| 74 | + try { |
| 75 | + $env:ATUIN_SHELL_POWERSHELL = "true" |
| 76 | + $argString = "search -i $ExtraArgs -- $line" |
| 77 | + Start-Process -Wait -NoNewWindow -RedirectStandardError $resultFile.FullName -FilePath atuin -ArgumentList $argString |
| 78 | + $suggestion = (Get-Content -Raw $resultFile -Encoding UTF8 | Out-String).Trim() |
| 79 | + } |
| 80 | + finally { |
| 81 | + $env:ATUIN_SHELL_POWERSHELL = $null |
| 82 | + Remove-Item $resultFile |
| 83 | + } |
| 84 | + |
| 85 | + $previousOutputEncoding = [System.Console]::OutputEncoding |
| 86 | + try { |
| 87 | + [System.Console]::OutputEncoding = [System.Text.Encoding]::UTF8 |
| 88 | + |
| 89 | + # PSReadLine maintains its own cursor position, which will no longer be valid if Atuin scrolls the display in inline mode. |
| 90 | + # Fortunately, InvokePrompt can receive a new Y position and reset the internal state. |
| 91 | + [Microsoft.PowerShell.PSConsoleReadLine]::InvokePrompt($null, $Host.UI.RawUI.CursorPosition.Y + [int]$env:ATUIN_POWERSHELL_PROMPT_OFFSET) |
| 92 | + |
| 93 | + if ($suggestion -eq "") { |
| 94 | + # The previous input was already rendered by InvokePrompt |
| 95 | + return |
| 96 | + } |
| 97 | + |
| 98 | + $acceptPrefix = "__atuin_accept__:" |
| 99 | + |
| 100 | + if ( $suggestion.StartsWith($acceptPrefix)) { |
| 101 | + [Microsoft.PowerShell.PSConsoleReadLine]::RevertLine() |
| 102 | + [Microsoft.PowerShell.PSConsoleReadLine]::Insert($suggestion.Substring($acceptPrefix.Length)) |
| 103 | + [Microsoft.PowerShell.PSConsoleReadLine]::AcceptLine() |
| 104 | + } else { |
| 105 | + [Microsoft.PowerShell.PSConsoleReadLine]::RevertLine() |
| 106 | + [Microsoft.PowerShell.PSConsoleReadLine]::Insert($suggestion) |
| 107 | + } |
| 108 | + } |
| 109 | + finally { |
| 110 | + [System.Console]::OutputEncoding = $previousOutputEncoding |
| 111 | + } |
| 112 | + } |
| 113 | + |
| 114 | + function Enable-AtuinSearchKeys { |
| 115 | + param([bool]$CtrlR = $true, [bool]$UpArrow = $true) |
| 116 | + |
| 117 | + if ($CtrlR) { |
| 118 | + Set-PSReadLineKeyHandler -Chord "Ctrl+r" -BriefDescription "Runs Atuin search" -ScriptBlock { |
| 119 | + RunSearch |
| 120 | + } |
| 121 | + } |
| 122 | + |
| 123 | + if ($UpArrow) { |
| 124 | + Set-PSReadLineKeyHandler -Chord "UpArrow" -BriefDescription "Runs Atuin search" -ScriptBlock { |
| 125 | + $line = $null |
| 126 | + [Microsoft.PowerShell.PSConsoleReadLine]::GetBufferState([ref]$line, [ref]$null) |
| 127 | + |
| 128 | + if (!$line.Contains("`n")) { |
| 129 | + RunSearch -ExtraArgs "--shell-up-key-binding" |
| 130 | + } else { |
| 131 | + [Microsoft.PowerShell.PSConsoleReadLine]::PreviousLine() |
| 132 | + } |
| 133 | + } |
| 134 | + } |
| 135 | + } |
| 136 | + |
| 137 | + $ExecutionContext.SessionState.Module.OnRemove += { |
| 138 | + $env:ATUIN_SESSION = $null |
| 139 | + $Function:PSConsoleHostReadLine = $script:previousPSConsoleHostReadLine |
| 140 | + } |
| 141 | + |
| 142 | + Export-ModuleMember -Function @("Enable-AtuinSearchKeys", "PSConsoleHostReadLine") |
| 143 | +} | Import-Module -Global |
0 commit comments