Skip to content

Add PowerShell module #2543

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions crates/atuin-common/src/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,11 @@ pub fn is_xonsh() -> bool {
env::var("ATUIN_SHELL_XONSH").is_ok()
}

pub fn is_powershell() -> bool {
// only set on powershell
env::var("ATUIN_SHELL_POWERSHELL").is_ok()
}

/// Extension trait for anything that can behave like a string to make it easy to escape control
/// characters.
///
Expand Down
1 change: 1 addition & 0 deletions crates/atuin-daemon/src/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ use atuin_client::encryption;
use atuin_client::history::store::HistoryStore;
use atuin_client::record::sqlite_store::SqliteStore;
use atuin_client::settings::Settings;
#[cfg(unix)]
use std::path::PathBuf;
use std::sync::Arc;
use time::OffsetDateTime;
Expand Down
1 change: 1 addition & 0 deletions crates/atuin-dotfiles/src/shell.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ use crate::store::AliasStore;

pub mod bash;
pub mod fish;
pub mod powershell;
pub mod xonsh;
pub mod zsh;

Expand Down
169 changes: 169 additions & 0 deletions crates/atuin-dotfiles/src/shell/powershell.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,169 @@
use crate::shell::{Alias, Var};
use crate::store::{AliasStore, var::VarStore};
use std::path::PathBuf;

async fn cached_aliases(path: PathBuf, store: &AliasStore) -> String {
match tokio::fs::read_to_string(path).await {
Ok(aliases) => aliases,
Err(r) => {
// we failed to read the file for some reason, but the file does exist
// fallback to generating new aliases on the fly

store.powershell().await.unwrap_or_else(|e| {
format!("echo 'Atuin: failed to read and generate aliases: \n{r}\n{e}'",)
})
}
}
}

async fn cached_vars(path: PathBuf, store: &VarStore) -> String {
match tokio::fs::read_to_string(path).await {
Ok(vars) => vars,
Err(r) => {
// we failed to read the file for some reason, but the file does exist
// fallback to generating new vars on the fly

store.powershell().await.unwrap_or_else(|e| {
format!("echo 'Atuin: failed to read and generate vars: \n{r}\n{e}'",)
})
}
}
}

/// Return powershell dotfile config
///
/// Do not return an error. We should not prevent the shell from starting.
///
/// In the worst case, Atuin should not function but the shell should start correctly.
///
/// While currently this only returns aliases, it will be extended to also return other synced dotfiles
pub async fn alias_config(store: &AliasStore) -> String {
// First try to read the cached config
let aliases = atuin_common::utils::dotfiles_cache_dir().join("aliases.ps1");

if aliases.exists() {
return cached_aliases(aliases, store).await;
}

if let Err(e) = store.build().await {
return format!("echo 'Atuin: failed to generate aliases: {}'", e);
}

cached_aliases(aliases, store).await
}

pub async fn var_config(store: &VarStore) -> String {
// First try to read the cached config
let vars = atuin_common::utils::dotfiles_cache_dir().join("vars.ps1");

if vars.exists() {
return cached_vars(vars, store).await;
}

if let Err(e) = store.build().await {
return format!("echo 'Atuin: failed to generate vars: {}'", e);
}

cached_vars(vars, store).await
}

pub fn format_alias(alias: &Alias) -> String {
// Set-Alias doesn't support adding implicit arguments, so use a function.
// See https://github.com/PowerShell/PowerShell/issues/12962

let mut result = secure_command(&format!(
"function {} {{\n {}{} @args\n}}",
alias.name,
if alias.value.starts_with(['"', '\'']) {
"& "
} else {
""
},
alias.value
));

// This makes the file layout prettier
result.insert(0, '\n');
result
}

pub fn format_var(var: &Var) -> String {
secure_command(&format!(
"${}{} = '{}'",
if var.export { "env:" } else { "" },
var.name,
var.value.replace("'", "''")
))
}

/// Wraps the given command in an Invoke-Expression to ensure the outer script is not halted
/// if the inner command contains a syntax error.
fn secure_command(command: &str) -> String {
format!(
"Invoke-Expression -ErrorAction Continue -Command '{}'\n",
command.replace("'", "''")
)
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn aliases() {
assert_eq!(
format_alias(&Alias {
name: "gp".to_string(),
value: "git push".to_string(),
}),
"\n".to_string()
+ &secure_command(
"function gp {
git push @args
}"
)
);

assert_eq!(
format_alias(&Alias {
name: "spc".to_string(),
value: "\"path with spaces\" arg".to_string(),
}),
"\n".to_string()
+ &secure_command(
"function spc {
& \"path with spaces\" arg @args
}"
)
);
}

#[test]
fn vars() {
assert_eq!(
format_var(&Var {
name: "FOO".to_owned(),
value: "bar 'baz'".to_owned(),
export: true,
}),
secure_command("$env:FOO = 'bar ''baz'''")
);

assert_eq!(
format_var(&Var {
name: "TEST".to_owned(),
value: "1".to_owned(),
export: false,
}),
secure_command("$TEST = '1'")
);
}

#[test]
fn invoke_expression() {
assert_eq!(
secure_command("echo 'foo'"),
"Invoke-Expression -ErrorAction Continue -Command 'echo ''foo'''\n"
)
}
}
40 changes: 33 additions & 7 deletions crates/atuin-dotfiles/src/store.rs
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,20 @@ impl AliasStore {

pub async fn posix(&self) -> Result<String> {
let aliases = self.aliases().await?;
Ok(Self::format_posix(&aliases))
}

pub async fn xonsh(&self) -> Result<String> {
let aliases = self.aliases().await?;
Ok(Self::format_xonsh(&aliases))
}

pub async fn powershell(&self) -> Result<String> {
let aliases = self.aliases().await?;
Ok(Self::format_powershell(&aliases))
}

fn format_posix(aliases: &[Alias]) -> String {
let mut config = String::new();

for alias in aliases {
Expand All @@ -153,28 +166,39 @@ impl AliasStore {
config.push_str(&format!("alias {}='{}'\n", alias.name, value));
}

Ok(config)
config
}

pub async fn xonsh(&self) -> Result<String> {
let aliases = self.aliases().await?;

fn format_xonsh(aliases: &[Alias]) -> String {
let mut config = String::new();

for alias in aliases {
config.push_str(&format!("aliases['{}'] ='{}'\n", alias.name, alias.value));
}

Ok(config)
config
}

fn format_powershell(aliases: &[Alias]) -> String {
let mut config = String::new();

for alias in aliases {
config.push_str(&crate::shell::powershell::format_alias(alias));
}

config
}

pub async fn build(&self) -> Result<()> {
let dir = atuin_common::utils::dotfiles_cache_dir();
tokio::fs::create_dir_all(dir.clone()).await?;

let aliases = self.aliases().await?;

// Build for all supported shells
let posix = self.posix().await?;
let xonsh = self.xonsh().await?;
let posix = Self::format_posix(&aliases);
let xonsh = Self::format_xonsh(&aliases);
let powershell = Self::format_powershell(&aliases);

// All the same contents, maybe optimize in the future or perhaps there will be quirks
// per-shell
Expand All @@ -183,11 +207,13 @@ impl AliasStore {
let bash = dir.join("aliases.bash");
let fish = dir.join("aliases.fish");
let xsh = dir.join("aliases.xsh");
let ps1 = dir.join("aliases.ps1");

tokio::fs::write(zsh, &posix).await?;
tokio::fs::write(bash, &posix).await?;
tokio::fs::write(fish, &posix).await?;
tokio::fs::write(xsh, &xonsh).await?;
tokio::fs::write(ps1, &powershell).await?;

Ok(())
}
Expand Down
53 changes: 41 additions & 12 deletions crates/atuin-dotfiles/src/store/var.rs
Original file line number Diff line number Diff line change
Expand Up @@ -117,31 +117,45 @@ impl VarStore {

pub async fn xonsh(&self) -> Result<String> {
let env = self.vars().await?;
Ok(Self::format_xonsh(&env))
}

pub async fn fish(&self) -> Result<String> {
let env = self.vars().await?;
Ok(Self::format_fish(&env))
}

pub async fn posix(&self) -> Result<String> {
let env = self.vars().await?;
Ok(Self::format_posix(&env))
}

pub async fn powershell(&self) -> Result<String> {
let env = self.vars().await?;
Ok(Self::format_powershell(&env))
}

fn format_xonsh(env: &[Var]) -> String {
let mut config = String::new();

for env in env {
config.push_str(&format!("${}={}\n", env.name, env.value));
}

Ok(config)
config
}

pub async fn fish(&self) -> Result<String> {
let env = self.vars().await?;

fn format_fish(env: &[Var]) -> String {
let mut config = String::new();

for env in env {
config.push_str(&format!("set -gx {} {}\n", env.name, env.value));
}

Ok(config)
config
}

pub async fn posix(&self) -> Result<String> {
let env = self.vars().await?;

fn format_posix(env: &[Var]) -> String {
let mut config = String::new();

for env in env {
Expand All @@ -152,17 +166,30 @@ impl VarStore {
}
}

Ok(config)
config
}

fn format_powershell(env: &[Var]) -> String {
let mut config = String::new();

for var in env {
config.push_str(&crate::shell::powershell::format_var(var));
}

config
}

pub async fn build(&self) -> Result<()> {
let dir = atuin_common::utils::dotfiles_cache_dir();
tokio::fs::create_dir_all(dir.clone()).await?;

let env = self.vars().await?;

// Build for all supported shells
let posix = self.posix().await?;
let xonsh = self.xonsh().await?;
let fsh = self.fish().await?;
let posix = Self::format_posix(&env);
let xonsh = Self::format_xonsh(&env);
let fsh = Self::format_fish(&env);
let powershell = Self::format_powershell(&env);

// All the same contents, maybe optimize in the future or perhaps there will be quirks
// per-shell
Expand All @@ -171,11 +198,13 @@ impl VarStore {
let bash = dir.join("vars.bash");
let fish = dir.join("vars.fish");
let xsh = dir.join("vars.xsh");
let ps1 = dir.join("vars.ps1");

tokio::fs::write(zsh, &posix).await?;
tokio::fs::write(bash, &posix).await?;
tokio::fs::write(fish, &fsh).await?;
tokio::fs::write(xsh, &xonsh).await?;
tokio::fs::write(ps1, &powershell).await?;

Ok(())
}
Expand Down
Loading
Loading