684 lines
20 KiB
Rust
684 lines
20 KiB
Rust
use std::env;
|
|
use std::io::Write;
|
|
use std::path::Path;
|
|
use std::process::{Command, Stdio};
|
|
use std::time::Duration;
|
|
use std::time::Instant;
|
|
|
|
use process_control::{ChildExt, Control, Output};
|
|
|
|
use super::{Context, Module, ModuleConfig};
|
|
|
|
use crate::{
|
|
config::Either, configs::custom::CustomConfig, formatter::StringFormatter,
|
|
utils::create_command,
|
|
};
|
|
|
|
/// Creates a custom module with some configuration
|
|
///
|
|
/// The relevant TOML config will set the files, extensions, and directories needed
|
|
/// for the module to be displayed. If none of them match, and optional "when"
|
|
/// command can be run -- if its result is 0, the module will be shown.
|
|
///
|
|
/// Finally, the content of the module itself is also set by a command.
|
|
pub fn module<'a>(name: &str, context: &'a Context) -> Option<Module<'a>> {
|
|
let start: Instant = Instant::now();
|
|
let toml_config = context.config.get_custom_module_config(name).expect(
|
|
"modules::custom::module should only be called after ensuring that the module exists",
|
|
);
|
|
let config = CustomConfig::load(toml_config);
|
|
|
|
if let Some(os) = config.os {
|
|
if os != env::consts::OS && !(os == "unix" && cfg!(unix)) {
|
|
return None;
|
|
}
|
|
}
|
|
|
|
let mut module = Module::new(name, config.description, Some(toml_config));
|
|
|
|
let mut is_match = context
|
|
.try_begin_scan()?
|
|
.set_extensions(&config.detect_extensions)
|
|
.set_files(&config.detect_files)
|
|
.set_folders(&config.detect_folders)
|
|
.is_match();
|
|
|
|
if !is_match {
|
|
is_match = match config.when {
|
|
Either::First(b) => b,
|
|
Either::Second(s) => exec_when(s, &config, context),
|
|
};
|
|
}
|
|
|
|
if is_match {
|
|
let parsed = StringFormatter::new(config.format).and_then(|formatter| {
|
|
formatter
|
|
.map_meta(|var, _| match var {
|
|
"symbol" => Some(config.symbol),
|
|
_ => None,
|
|
})
|
|
.map_style(|variable| match variable {
|
|
"style" => Some(Ok(config.style)),
|
|
_ => None,
|
|
})
|
|
.map_no_escaping(|variable| match variable {
|
|
"output" => {
|
|
let output = exec_command(config.command, context, &config)?;
|
|
let trimmed = output.trim();
|
|
|
|
if trimmed.is_empty() {
|
|
None
|
|
} else {
|
|
Some(Ok(trimmed.to_string()))
|
|
}
|
|
}
|
|
_ => None,
|
|
})
|
|
.parse(None, Some(context))
|
|
});
|
|
|
|
match parsed {
|
|
Ok(segments) => module.set_segments(segments),
|
|
Err(error) => {
|
|
log::warn!("Error in module `custom.{}`:\n{}", name, error);
|
|
}
|
|
};
|
|
}
|
|
let elapsed = start.elapsed();
|
|
log::trace!("Took {:?} to compute custom module {:?}", elapsed, name);
|
|
module.duration = elapsed;
|
|
Some(module)
|
|
}
|
|
|
|
/// Return the invoking shell, using `shell` and fallbacking in order to `STARSHIP_SHELL` and "sh"/"cmd"
|
|
fn get_shell<'a, 'b>(
|
|
shell_args: &'b [&'a str],
|
|
context: &Context,
|
|
) -> (std::borrow::Cow<'a, str>, &'b [&'a str]) {
|
|
if !shell_args.is_empty() {
|
|
(shell_args[0].into(), &shell_args[1..])
|
|
} else if let Some(env_shell) = context.get_env("STARSHIP_SHELL") {
|
|
(env_shell.into(), &[] as &[&str])
|
|
} else if cfg!(windows) {
|
|
// `/C` is added by `handle_shell`
|
|
("cmd".into(), &[] as &[&str])
|
|
} else {
|
|
("sh".into(), &[] as &[&str])
|
|
}
|
|
}
|
|
|
|
/// Attempt to run the given command in a shell by passing it as either `stdin` or an argument to `get_shell()`,
|
|
/// depending on the configuration or by invoking a platform-specific fallback shell if `shell` is empty.
|
|
fn shell_command(cmd: &str, config: &CustomConfig, context: &Context) -> Option<Output> {
|
|
let (shell, shell_args) = get_shell(config.shell.0.as_ref(), context);
|
|
let mut use_stdin = config.use_stdin;
|
|
|
|
let mut command = match create_command(shell.as_ref()) {
|
|
Ok(command) => command,
|
|
// Don't attempt to use fallback shell if the user specified a shell
|
|
Err(error) if !shell_args.is_empty() => {
|
|
log::debug!(
|
|
"Error creating command with STARSHIP_SHELL, falling back to fallback shell: {}",
|
|
error
|
|
);
|
|
|
|
// Skip `handle_shell` and just set the shell and command
|
|
use_stdin = Some(!cfg!(windows));
|
|
|
|
if cfg!(windows) {
|
|
let mut c = create_command("cmd").ok()?;
|
|
c.arg("/C");
|
|
c
|
|
} else {
|
|
let mut c = create_command("/usr/bin/env").ok()?;
|
|
c.arg("sh");
|
|
c
|
|
}
|
|
}
|
|
_ => return None,
|
|
};
|
|
|
|
command
|
|
.current_dir(&context.current_dir)
|
|
.args(shell_args)
|
|
.stdin(Stdio::piped())
|
|
.stdout(Stdio::piped())
|
|
.stderr(Stdio::piped());
|
|
|
|
let use_stdin = use_stdin.unwrap_or_else(|| handle_shell(&mut command, &shell, shell_args));
|
|
|
|
if !use_stdin {
|
|
command.arg(cmd);
|
|
}
|
|
|
|
let mut child = match command.spawn() {
|
|
Ok(child) => child,
|
|
Err(error) => {
|
|
log::debug!(
|
|
"Failed to run command with given shell or STARSHIP_SHELL env variable:: {}",
|
|
error
|
|
);
|
|
return None;
|
|
}
|
|
};
|
|
|
|
if use_stdin {
|
|
child.stdin.as_mut()?.write_all(cmd.as_bytes()).ok()?;
|
|
}
|
|
|
|
let mut output = child.controlled_with_output();
|
|
|
|
if !config.ignore_timeout {
|
|
output = output
|
|
.time_limit(Duration::from_millis(context.root_config.command_timeout))
|
|
.terminate_for_timeout()
|
|
}
|
|
|
|
match output.wait().ok()? {
|
|
None => {
|
|
log::warn!("Executing custom command {cmd:?} timed out.");
|
|
log::warn!("You can set command_timeout in your config to a higher value or set ignore_timeout to true for this module to allow longer-running commands to keep executing.");
|
|
None
|
|
}
|
|
Some(status) => Some(status),
|
|
}
|
|
}
|
|
|
|
/// Execute the given command capturing all output, and return whether it return 0
|
|
fn exec_when(cmd: &str, config: &CustomConfig, context: &Context) -> bool {
|
|
log::trace!("Running '{}'", cmd);
|
|
|
|
if let Some(output) = shell_command(cmd, config, context) {
|
|
if !output.status.success() {
|
|
log::trace!("non-zero exit code '{:?}'", output.status.code());
|
|
log::trace!(
|
|
"stdout: {}",
|
|
std::str::from_utf8(&output.stdout).unwrap_or("<invalid utf8>")
|
|
);
|
|
log::trace!(
|
|
"stderr: {}",
|
|
std::str::from_utf8(&output.stderr).unwrap_or("<invalid utf8>")
|
|
);
|
|
}
|
|
|
|
output.status.success()
|
|
} else {
|
|
log::debug!("Cannot start command");
|
|
|
|
false
|
|
}
|
|
}
|
|
|
|
/// Execute the given command, returning its output on success
|
|
fn exec_command(cmd: &str, context: &Context, config: &CustomConfig) -> Option<String> {
|
|
log::trace!("Running '{cmd}'");
|
|
|
|
if let Some(output) = shell_command(cmd, config, context) {
|
|
if !output.status.success() {
|
|
log::trace!("Non-zero exit code '{:?}'", output.status.code());
|
|
log::trace!(
|
|
"stdout: {}",
|
|
std::str::from_utf8(&output.stdout).unwrap_or("<invalid utf8>")
|
|
);
|
|
log::trace!(
|
|
"stderr: {}",
|
|
std::str::from_utf8(&output.stderr).unwrap_or("<invalid utf8>")
|
|
);
|
|
return None;
|
|
}
|
|
|
|
Some(String::from_utf8_lossy(&output.stdout).into())
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
/// If the specified shell refers to `PowerShell`, adds the arguments "-Command -" to the
|
|
/// given command.
|
|
/// Returns `false` if the shell shell expects scripts as arguments, `true` if as `stdin`.
|
|
fn handle_shell(command: &mut Command, shell: &str, shell_args: &[&str]) -> bool {
|
|
let shell_exe = Path::new(shell).file_stem();
|
|
let no_args = shell_args.is_empty();
|
|
|
|
match shell_exe.and_then(std::ffi::OsStr::to_str) {
|
|
Some("pwsh" | "powershell") => {
|
|
if no_args {
|
|
command.arg("-NoProfile").arg("-Command").arg("-");
|
|
}
|
|
true
|
|
}
|
|
Some("cmd") => {
|
|
if no_args {
|
|
command.arg("/C");
|
|
}
|
|
false
|
|
}
|
|
Some("nu") => {
|
|
if no_args {
|
|
command.arg("-c");
|
|
}
|
|
false
|
|
}
|
|
_ => true,
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
use crate::test::ModuleRenderer;
|
|
use nu_ansi_term::Color;
|
|
use std::fs::File;
|
|
use std::io;
|
|
|
|
#[cfg(not(windows))]
|
|
const SHELL: &[&str] = &["/bin/sh"];
|
|
#[cfg(windows)]
|
|
const SHELL: &[&str] = &["cmd"];
|
|
|
|
#[cfg(not(windows))]
|
|
const FAILING_COMMAND: &str = "false";
|
|
#[cfg(windows)]
|
|
const FAILING_COMMAND: &str = "color 00";
|
|
|
|
const UNKNOWN_COMMAND: &str = "ydelsyiedsieudleylse dyesdesl";
|
|
|
|
fn render_cmd(cmd: &str) -> io::Result<Option<String>> {
|
|
let dir = tempfile::tempdir()?;
|
|
let cmd = cmd.to_owned();
|
|
let shell = SHELL
|
|
.iter()
|
|
.map(std::borrow::ToOwned::to_owned)
|
|
.collect::<Vec<_>>();
|
|
let out = ModuleRenderer::new("custom.test")
|
|
.path(dir.path())
|
|
.config(toml::toml! {
|
|
[custom.test]
|
|
format = "$output"
|
|
command = cmd
|
|
shell = shell
|
|
when = true
|
|
ignore_timeout = true
|
|
})
|
|
.collect();
|
|
dir.close()?;
|
|
Ok(out)
|
|
}
|
|
|
|
fn render_when(cmd: &str) -> io::Result<bool> {
|
|
let dir = tempfile::tempdir()?;
|
|
let cmd = cmd.to_owned();
|
|
let shell = SHELL
|
|
.iter()
|
|
.map(std::borrow::ToOwned::to_owned)
|
|
.collect::<Vec<_>>();
|
|
let out = ModuleRenderer::new("custom.test")
|
|
.path(dir.path())
|
|
.config(toml::toml! {
|
|
[custom.test]
|
|
format = "test"
|
|
when = cmd
|
|
shell = shell
|
|
ignore_timeout = true
|
|
})
|
|
.collect()
|
|
.is_some();
|
|
dir.close()?;
|
|
Ok(out)
|
|
}
|
|
|
|
#[test]
|
|
fn when_returns_right_value() -> io::Result<()> {
|
|
assert!(render_cmd("echo hello")?.is_some());
|
|
assert!(render_cmd(FAILING_COMMAND)?.is_none());
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn when_returns_false_if_invalid_command() -> io::Result<()> {
|
|
assert!(!render_when(UNKNOWN_COMMAND)?);
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(not(windows))]
|
|
fn command_returns_right_string() -> io::Result<()> {
|
|
assert_eq!(render_cmd("echo hello")?, Some("hello".into()));
|
|
assert_eq!(render_cmd("echo 강남스타일")?, Some("강남스타일".into()));
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(windows)]
|
|
fn command_returns_right_string() -> io::Result<()> {
|
|
assert_eq!(render_cmd("echo hello")?, Some("hello".into()));
|
|
assert_eq!(render_cmd("echo 강남스타일")?, Some("강남스타일".into()));
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(not(windows))]
|
|
fn command_ignores_stderr() -> io::Result<()> {
|
|
assert_eq!(render_cmd("echo foo 1>&2; echo bar")?, Some("bar".into()));
|
|
assert_eq!(render_cmd("echo foo; echo bar 1>&2")?, Some("foo".into()));
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(windows)]
|
|
fn command_ignores_stderr() -> io::Result<()> {
|
|
assert_eq!(render_cmd("echo foo 1>&2 & echo bar")?, Some("bar".into()));
|
|
assert_eq!(render_cmd("echo foo& echo bar 1>&2")?, Some("foo".into()));
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn command_can_fail() -> io::Result<()> {
|
|
assert_eq!(render_cmd(FAILING_COMMAND)?, None);
|
|
assert_eq!(render_cmd(UNKNOWN_COMMAND)?, None);
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn cwd_command() -> io::Result<()> {
|
|
let dir = tempfile::tempdir()?;
|
|
|
|
let mut f = File::create(dir.path().join("a.txt"))?;
|
|
write!(f, "hello")?;
|
|
f.sync_all()?;
|
|
|
|
let cat = if cfg!(windows) { "type" } else { "cat" };
|
|
let cmd = format!("{cat} a.txt");
|
|
|
|
let actual = ModuleRenderer::new("custom.test")
|
|
.path(dir.path())
|
|
.config(toml::toml! {
|
|
[custom.test]
|
|
command = cmd
|
|
when = true
|
|
ignore_timeout = true
|
|
})
|
|
.collect();
|
|
let expected = Some(format!("{}", Color::Green.bold().paint("hello ")));
|
|
|
|
assert_eq!(expected, actual);
|
|
|
|
dir.close()
|
|
}
|
|
|
|
#[test]
|
|
fn cwd_when() -> io::Result<()> {
|
|
let dir = tempfile::tempdir()?;
|
|
|
|
File::create(dir.path().join("a.txt"))?.sync_all()?;
|
|
|
|
let cat = if cfg!(windows) { "type" } else { "cat" };
|
|
let cmd = format!("{cat} a.txt");
|
|
|
|
let actual = ModuleRenderer::new("custom.test")
|
|
.path(dir.path())
|
|
.config(toml::toml! {
|
|
[custom.test]
|
|
format = "test"
|
|
when = cmd
|
|
ignore_timeout = true
|
|
})
|
|
.collect();
|
|
let expected = Some("test".to_owned());
|
|
|
|
assert_eq!(expected, actual);
|
|
|
|
dir.close()
|
|
}
|
|
|
|
#[test]
|
|
fn use_stdin_false() -> io::Result<()> {
|
|
let dir = tempfile::tempdir()?;
|
|
|
|
let shell = if cfg!(windows) {
|
|
vec![
|
|
"powershell".to_owned(),
|
|
"-NoProfile".to_owned(),
|
|
"-Command".to_owned(),
|
|
]
|
|
} else {
|
|
vec!["sh".to_owned(), "-c".to_owned()]
|
|
};
|
|
|
|
// `use_stdin = false` doesn't like Korean on Windows
|
|
let actual = ModuleRenderer::new("custom.test")
|
|
.path(dir.path())
|
|
.config(toml::toml! {
|
|
[custom.test]
|
|
command = "echo test"
|
|
when = true
|
|
use_stdin = false
|
|
shell = shell
|
|
ignore_timeout = true
|
|
})
|
|
.collect();
|
|
let expected = Some(format!("{}", Color::Green.bold().paint("test ")));
|
|
|
|
assert_eq!(expected, actual);
|
|
|
|
dir.close()
|
|
}
|
|
|
|
#[test]
|
|
fn use_stdin_true() -> io::Result<()> {
|
|
let dir = tempfile::tempdir()?;
|
|
|
|
let shell = if cfg!(windows) {
|
|
vec![
|
|
"powershell".to_owned(),
|
|
"-NoProfile".to_owned(),
|
|
"-Command".to_owned(),
|
|
"-".to_owned(),
|
|
]
|
|
} else {
|
|
vec!["sh".to_owned()]
|
|
};
|
|
|
|
let actual = ModuleRenderer::new("custom.test")
|
|
.path(dir.path())
|
|
.config(toml::toml! {
|
|
[custom.test]
|
|
command = "echo 강남스타일"
|
|
when = true
|
|
use_stdin = true
|
|
ignore_timeout = true
|
|
shell = shell
|
|
})
|
|
.collect();
|
|
let expected = Some(format!("{}", Color::Green.bold().paint("강남스타일 ")));
|
|
|
|
assert_eq!(expected, actual);
|
|
|
|
dir.close()
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(not(windows))]
|
|
fn when_true_with_string() -> std::io::Result<()> {
|
|
let dir = tempfile::tempdir()?;
|
|
|
|
let actual = ModuleRenderer::new("custom.test")
|
|
.path(dir.path())
|
|
.config(toml::toml! {
|
|
[custom.test]
|
|
format = "test"
|
|
shell = ["sh"]
|
|
when = "true"
|
|
ignore_timeout = true
|
|
})
|
|
.collect();
|
|
let expected = Some("test".to_string());
|
|
assert_eq!(expected, actual);
|
|
|
|
dir.close()
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(not(windows))]
|
|
fn when_false_with_string() -> std::io::Result<()> {
|
|
let dir = tempfile::tempdir()?;
|
|
|
|
let actual = ModuleRenderer::new("custom.test")
|
|
.path(dir.path())
|
|
.config(toml::toml! {
|
|
[custom.test]
|
|
format = "test"
|
|
shell = ["sh"]
|
|
when = "false"
|
|
ignore_timeout = true
|
|
})
|
|
.collect();
|
|
let expected = None;
|
|
assert_eq!(expected, actual);
|
|
|
|
dir.close()
|
|
}
|
|
|
|
#[test]
|
|
fn when_true_with_bool() -> std::io::Result<()> {
|
|
let dir = tempfile::tempdir()?;
|
|
|
|
let actual = ModuleRenderer::new("custom.test")
|
|
.path(dir.path())
|
|
.config(toml::toml! {
|
|
[custom.test]
|
|
format = "test"
|
|
when = true
|
|
})
|
|
.collect();
|
|
let expected = Some("test".to_string());
|
|
assert_eq!(expected, actual);
|
|
|
|
dir.close()
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(not(windows))]
|
|
fn when_false_with_bool() -> std::io::Result<()> {
|
|
let dir = tempfile::tempdir()?;
|
|
|
|
let actual = ModuleRenderer::new("custom.test")
|
|
.path(dir.path())
|
|
.config(toml::toml! {
|
|
[custom.test]
|
|
format = "test"
|
|
when = false
|
|
})
|
|
.collect();
|
|
let expected = None;
|
|
assert_eq!(expected, actual);
|
|
|
|
dir.close()
|
|
}
|
|
|
|
#[test]
|
|
fn timeout_short_cmd() -> std::io::Result<()> {
|
|
let dir = tempfile::tempdir()?;
|
|
|
|
let shell = if cfg!(windows) {
|
|
"powershell".to_owned()
|
|
} else {
|
|
"sh".to_owned()
|
|
};
|
|
|
|
let when = if cfg!(windows) {
|
|
"$true".to_owned()
|
|
} else {
|
|
"true".to_owned()
|
|
};
|
|
|
|
// Use a long timeout to ensure that the test doesn't fail
|
|
let actual = ModuleRenderer::new("custom.test")
|
|
.path(dir.path())
|
|
.config(toml::toml! {
|
|
command_timeout = 100_000
|
|
[custom.test]
|
|
format = "test"
|
|
when = when
|
|
shell = shell
|
|
ignore_timeout = false
|
|
})
|
|
.collect();
|
|
let expected = Some("test".to_owned());
|
|
assert_eq!(expected, actual);
|
|
|
|
dir.close()
|
|
}
|
|
|
|
#[test]
|
|
fn timeout_cmd() -> std::io::Result<()> {
|
|
let dir = tempfile::tempdir()?;
|
|
|
|
let shell = if cfg!(windows) {
|
|
"powershell".to_owned()
|
|
} else {
|
|
"sh".to_owned()
|
|
};
|
|
|
|
// Use a long timeout to ensure that the test doesn't fail
|
|
let actual = ModuleRenderer::new("custom.test")
|
|
.path(dir.path())
|
|
.config(toml::toml! {
|
|
[custom.test]
|
|
format = "test"
|
|
when = "sleep 3"
|
|
shell = shell
|
|
ignore_timeout = false
|
|
})
|
|
.collect();
|
|
let expected = None;
|
|
assert_eq!(expected, actual);
|
|
|
|
dir.close()
|
|
}
|
|
|
|
#[test]
|
|
fn config_aliases_work() -> std::io::Result<()> {
|
|
let dir = tempfile::tempdir()?;
|
|
|
|
File::create(dir.path().join("a.txt"))?;
|
|
std::fs::create_dir(dir.path().join("dir"))?;
|
|
|
|
let actual = ModuleRenderer::new("custom.test")
|
|
.path(dir.path())
|
|
.config(toml::toml! {
|
|
[custom.test]
|
|
format = "test"
|
|
files = ["a.txt"]
|
|
})
|
|
.collect();
|
|
let expected = Some("test".to_string());
|
|
assert_eq!(expected, actual);
|
|
|
|
let actual = ModuleRenderer::new("custom.test")
|
|
.path(dir.path())
|
|
.config(toml::toml! {
|
|
[custom.test]
|
|
format = "test"
|
|
extensions = ["txt"]
|
|
})
|
|
.collect();
|
|
let expected = Some("test".to_string());
|
|
assert_eq!(expected, actual);
|
|
|
|
let actual = ModuleRenderer::new("custom.test")
|
|
.path(dir.path())
|
|
.config(toml::toml! {
|
|
[custom.test]
|
|
format = "test"
|
|
directories = ["dir"]
|
|
})
|
|
.collect();
|
|
let expected = Some("test".to_string());
|
|
assert_eq!(expected, actual);
|
|
|
|
dir.close()
|
|
}
|
|
}
|