Skip to content

Commit 567ee37

Browse files
ltrzesniewskilzybkr
ltrzesniewski
and
lzybkr
committed
feat: add PowerShell module
This adds PowerShell support by invoking the following expression: atuin init powershell | Out-String | Invoke-Expression Co-authored-by: Jason Shirk <[email protected]>
1 parent 05aec6f commit 567ee37

File tree

5 files changed

+214
-1
lines changed

5 files changed

+214
-1
lines changed

‎crates/atuin-common/src/utils.rs

+5
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,11 @@ pub fn is_xonsh() -> bool {
133133
env::var("ATUIN_SHELL_XONSH").is_ok()
134134
}
135135

136+
pub fn is_powershell() -> bool {
137+
// only set on powershell
138+
env::var("ATUIN_SHELL_POWERSHELL").is_ok()
139+
}
140+
136141
/// Extension trait for anything that can behave like a string to make it easy to escape control
137142
/// characters.
138143
///

‎crates/atuin/src/command/client/init.rs

+16
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use eyre::{Result, WrapErr};
77

88
mod bash;
99
mod fish;
10+
mod powershell;
1011
mod xonsh;
1112
mod zsh;
1213

@@ -24,6 +25,7 @@ pub struct Cmd {
2425
}
2526

2627
#[derive(Clone, Copy, ValueEnum, Debug)]
28+
#[value(rename_all = "lower")]
2729
pub enum Shell {
2830
/// Zsh setup
2931
Zsh,
@@ -35,6 +37,8 @@ pub enum Shell {
3537
Nu,
3638
/// Xonsh setup
3739
Xonsh,
40+
/// PowerShell setup
41+
PowerShell,
3842
}
3943

4044
impl Cmd {
@@ -100,6 +104,9 @@ $env.config = (
100104
Shell::Xonsh => {
101105
xonsh::init_static(self.disable_up_arrow, self.disable_ctrl_r);
102106
}
107+
Shell::PowerShell => {
108+
powershell::init_static(self.disable_up_arrow, self.disable_ctrl_r);
109+
}
103110
};
104111
}
105112

@@ -153,6 +160,15 @@ $env.config = (
153160
)
154161
.await?;
155162
}
163+
Shell::PowerShell => {
164+
powershell::init(
165+
alias_store,
166+
var_store,
167+
self.disable_up_arrow,
168+
self.disable_ctrl_r,
169+
)
170+
.await?;
171+
}
156172
}
157173

158174
Ok(())

‎crates/atuin/src/command/client/init/powershell.rs

+40
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
use atuin_dotfiles::store::{var::VarStore, AliasStore};
2+
use eyre::Result;
3+
4+
pub fn init_static(disable_up_arrow: bool, disable_ctrl_r: bool) {
5+
let base = include_str!("../../../shell/atuin.ps1");
6+
7+
let (bind_ctrl_r, bind_up_arrow) = if std::env::var("ATUIN_NOBIND").is_ok() {
8+
(false, false)
9+
} else {
10+
(!disable_ctrl_r, !disable_up_arrow)
11+
};
12+
13+
fn bool(value: bool) -> &'static str {
14+
if value {
15+
"$true"
16+
} else {
17+
"$false"
18+
}
19+
}
20+
21+
println!("{base}");
22+
println!(
23+
"Enable-AtuinSearchKeys -CtrlR {} -UpArrow {}",
24+
bool(bind_ctrl_r),
25+
bool(bind_up_arrow)
26+
);
27+
}
28+
29+
pub async fn init(
30+
_aliases: AliasStore,
31+
_vars: VarStore,
32+
disable_up_arrow: bool,
33+
disable_ctrl_r: bool,
34+
) -> Result<()> {
35+
init_static(disable_up_arrow, disable_ctrl_r);
36+
37+
// dotfiles are not supported yet
38+
39+
Ok(())
40+
}

‎crates/atuin/src/command/client/search/interactive.rs

+9-1
Original file line numberDiff line numberDiff line change
@@ -1098,6 +1098,10 @@ pub async fn history(
10981098

10991099
let mut results = app.query_results(&mut db, settings.smart_sort).await?;
11001100

1101+
if settings.inline_height > 0 {
1102+
terminal.clear()?;
1103+
}
1104+
11011105
let mut stats: Option<HistoryStats> = None;
11021106
let accept;
11031107
let result = 'render: loop {
@@ -1180,7 +1184,11 @@ pub async fn history(
11801184
InputAction::Accept(index) if index < results.len() => {
11811185
let mut command = results.swap_remove(index).command;
11821186
if accept
1183-
&& (utils::is_zsh() || utils::is_fish() || utils::is_bash() || utils::is_xonsh())
1187+
&& (utils::is_zsh()
1188+
|| utils::is_fish()
1189+
|| utils::is_bash()
1190+
|| utils::is_xonsh()
1191+
|| utils::is_powershell())
11841192
{
11851193
command = String::from("__atuin_accept__:") + &command;
11861194
}

‎crates/atuin/src/shell/atuin.ps1

+144
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
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+
$cursor = $null
68+
[Microsoft.PowerShell.PSConsoleReadLine]::GetBufferState([ref]$line, [ref]$cursor)
69+
70+
# Atuin is started through Start-Process to avoid interfering with the current shell,
71+
# and to capture its output which is provided in stderr (redirected to a temporary file).
72+
73+
$suggestion = ""
74+
$resultFile = New-TemporaryFile
75+
try {
76+
$env:ATUIN_SHELL_POWERSHELL = "true"
77+
$argString = "search -i $ExtraArgs -- $line"
78+
Start-Process -Wait -NoNewWindow -RedirectStandardError $resultFile.FullName -FilePath atuin -ArgumentList $argString
79+
$suggestion = (Get-Content -Raw $resultFile | Out-String).Trim()
80+
}
81+
finally {
82+
$env:ATUIN_SHELL_POWERSHELL = $null
83+
Remove-Item $resultFile
84+
}
85+
86+
$previousOutputEncoding = [System.Console]::OutputEncoding
87+
try {
88+
[System.Console]::OutputEncoding = [System.Text.Encoding]::UTF8
89+
90+
# PSReadLine maintains its own cursor position, which will no longer be valid if Atuin scrolls the display in inline mode.
91+
# Fortunately, InvokePrompt can receive a new Y position and reset the internal state.
92+
[Microsoft.PowerShell.PSConsoleReadLine]::InvokePrompt($null, $Host.UI.RawUI.CursorPosition.Y)
93+
}
94+
finally {
95+
[System.Console]::OutputEncoding = $previousOutputEncoding
96+
}
97+
98+
if ($suggestion -eq "") {
99+
# The previous input was already rendered by InvokePrompt
100+
return
101+
}
102+
103+
$acceptPrefix = "__atuin_accept__:"
104+
105+
if ( $suggestion.StartsWith($acceptPrefix)) {
106+
[Microsoft.PowerShell.PSConsoleReadLine]::RevertLine()
107+
[Microsoft.PowerShell.PSConsoleReadLine]::Insert($suggestion.Substring($acceptPrefix.Length))
108+
[Microsoft.PowerShell.PSConsoleReadLine]::AcceptLine()
109+
} else {
110+
[Microsoft.PowerShell.PSConsoleReadLine]::RevertLine()
111+
[Microsoft.PowerShell.PSConsoleReadLine]::Insert($suggestion)
112+
}
113+
}
114+
115+
function Enable-AtuinSearchKeys {
116+
param([bool]$CtrlR = $true, [bool]$UpArrow = $true)
117+
118+
if ($CtrlR) {
119+
Set-PSReadLineKeyHandler -Chord "Ctrl+r" -BriefDescription "Runs Atuin search" -ScriptBlock {
120+
RunSearch
121+
}
122+
}
123+
124+
if ($UpArrow) {
125+
Set-PSReadLineKeyHandler -Chord "UpArrow" -BriefDescription "Runs Atuin search" -ScriptBlock {
126+
$line = $null
127+
[Microsoft.PowerShell.PSConsoleReadLine]::GetBufferState([ref]$line, [ref]$null)
128+
129+
if (!$line.Contains("`n")) {
130+
RunSearch -ExtraArgs "--shell-up-key-binding"
131+
} else {
132+
[Microsoft.PowerShell.PSConsoleReadLine]::PreviousLine()
133+
}
134+
}
135+
}
136+
}
137+
138+
$ExecutionContext.SessionState.Module.OnRemove += {
139+
$env:ATUIN_SESSION = $null
140+
$Function:PSConsoleHostReadLine = $script:previousPSConsoleHostReadLine
141+
}
142+
143+
Export-ModuleMember -Function @("Enable-AtuinSearchKeys", "PSConsoleHostReadLine")
144+
} | Import-Module -Global

0 commit comments

Comments
 (0)