use std::env;
#[cfg(target_os = "linux")]
use std::fs;
use std::io::Read;
use std::process::ExitCode;
use std::str::FromStr;
use std::sync::Arc;
use std::time::{Duration, Instant};
mod detector;
mod error;
mod executor;
mod memo;
mod module_trait;
mod modules;
mod parser;
mod registry;
mod style;
const VERSION: &str = env!("CARGO_PKG_VERSION");
const HELP: &str = "\
prmt - Ultra-fast customizable shell prompt generator
USAGE:
prmt [OPTIONS] [FORMAT]
ARGS:
<FORMAT> Format string (default from PRMT_FORMAT env var)
OPTIONS:
-f, --format <FORMAT> Format string
-n, --no-version Skip version detection for speed
-t, --timeout <MS> Prompt timeout in ms (default 0, disabled)
-d, --debug Show debug information and timing
-b, --bench Run benchmark (100 iterations, ignores module timeout)
--stdin Read JSON from stdin (enables json module)
--code <CODE> Exit code of the last command (for ok/fail modules)
--no-color Disable colored output
--shell <SHELL> Wrap ANSI escapes for the specified shell (bash, zsh, none)
-h, --help Print help
-V, --version Print version
";
struct Cli {
format: Option<String>,
no_version: bool,
timeout: Option<u64>,
debug: bool,
bench: bool,
stdin: bool,
code: Option<i32>,
no_color: bool,
shell: Option<style::Shell>,
}
fn parse_args() -> Result<Cli, lexopt::Error> {
parse_args_from(std::env::args_os())
}
fn parse_args_from<I, T>(args: I) -> Result<Cli, lexopt::Error>
where
I: IntoIterator<Item = T>,
T: Into<std::ffi::OsString>,
{
use lexopt::prelude::*;
let mut format = None;
let mut no_version = false;
let mut timeout = None;
let mut debug = false;
let mut bench = false;
let mut stdin = false;
let mut code = None;
let mut no_color = false;
let mut shell = None;
let mut parser = lexopt::Parser::from_iter(args);
while let Some(arg) = parser.next()? {
match arg {
Short('h') | Long("help") => {
print!("{}", HELP);
std::process::exit(0);
}
Short('V') | Long("version") => {
println!("prmt {}", VERSION);
std::process::exit(0);
}
Short('f') | Long("format") => {
format = Some(parser.value()?.string()?);
}
Short('n') | Long("no-version") => {
no_version = true;
}
Short('t') | Long("timeout") => {
timeout = Some(parser.value()?.parse()?);
}
Short('d') | Long("debug") => {
debug = true;
}
Short('b') | Long("bench") => {
bench = true;
}
Long("stdin") => {
stdin = true;
}
Long("code") => {
code = Some(parser.value()?.parse()?);
}
Long("no-color") => {
no_color = true;
}
Long("shell") => {
let value = parser.value()?.string()?;
shell = Some(style::Shell::from_str(&value)?);
}
arg if matches!(&arg, Value(_)) => {
let Value(val) = arg else { unreachable!() };
if format.is_none() {
format = Some(val.string()?);
} else {
return Err(Value(val).unexpected());
}
}
_ => return Err(arg.unexpected()),
}
}
Ok(Cli {
format,
no_version,
timeout,
debug,
bench,
stdin,
code,
no_color,
shell,
})
}
fn shell_from_name(value: &str) -> Option<style::Shell> {
let trimmed = value.trim().trim_end_matches('\0').trim_start_matches('-');
let name = trimmed.rsplit('/').next().unwrap_or(trimmed);
match name {
"zsh" => Some(style::Shell::Zsh),
"bash" => Some(style::Shell::Bash),
_ => None,
}
}
fn detect_shell_from_env() -> Option<style::Shell> {
if env::var("ZSH_VERSION").is_ok() {
return Some(style::Shell::Zsh);
}
if env::var("BASH_VERSION").is_ok() {
return Some(style::Shell::Bash);
}
if let Ok(shell_path) = env::var("SHELL")
&& let Some(shell) = shell_from_name(&shell_path)
{
return Some(shell);
}
None
}
#[cfg(target_os = "linux")]
fn detect_shell_from_parent_process() -> Option<style::Shell> {
let status = fs::read_to_string("/proc/self/status").ok()?;
let ppid_line = status.lines().find(|line| line.starts_with("PPid:"))?;
let ppid = ppid_line.split_whitespace().nth(1)?.parse::<u32>().ok()?;
if ppid == 0 {
return None;
}
let comm = fs::read_to_string(format!("/proc/{}/comm", ppid)).ok()?;
if let Some(shell) = shell_from_name(&comm) {
return Some(shell);
}
let cmdline = fs::read_to_string(format!("/proc/{}/cmdline", ppid)).ok()?;
let first = cmdline.split('\0').next().unwrap_or("");
shell_from_name(first)
}
#[cfg(not(target_os = "linux"))]
fn detect_shell_from_parent_process() -> Option<style::Shell> {
None
}
fn resolve_shell(cli_shell: Option<style::Shell>) -> style::Shell {
resolve_shell_from_sources(
cli_shell,
detect_shell_from_parent_process(),
detect_shell_from_env(),
)
}
fn resolve_shell_from_sources(
cli_shell: Option<style::Shell>,
parent_shell: Option<style::Shell>,
env_shell: Option<style::Shell>,
) -> style::Shell {
if let Some(shell) = cli_shell {
return shell;
}
if let Some(shell) = parent_shell {
return shell;
}
env_shell.unwrap_or(style::Shell::None)
}
fn main() -> ExitCode {
let cli = match parse_args() {
Ok(cli) => cli,
Err(e) => {
eprintln!("Error: {}", e);
eprintln!("Try 'prmt --help' for more information.");
return ExitCode::FAILURE;
}
};
let format = cli
.format
.or_else(|| env::var("PRMT_FORMAT").ok())
.unwrap_or_else(|| "{path:cyan} {node:green} {git:purple}".to_string());
let shell = resolve_shell(cli.shell);
let timeout_ms = cli
.timeout
.or_else(|| env::var("PRMT_TIMEOUT").ok()?.parse().ok())
.unwrap_or(0);
let timeout = if timeout_ms > 0 {
Some(Duration::from_millis(timeout_ms))
} else {
None
};
let stdin_data = if cli.stdin { read_stdin_json() } else { None };
let result = if cli.bench {
handle_bench(
&format,
cli.no_version,
cli.code,
cli.no_color,
shell,
stdin_data,
)
} else {
let start = cli.debug.then(Instant::now);
let output = executor::execute_with_shell(
&format,
cli.no_version,
cli.code,
cli.no_color,
shell,
stdin_data,
timeout,
);
if let Some(start) = start {
let elapsed = start.elapsed();
eprintln!("Format: {}", format);
eprintln!("Execution time: {:.2}ms", elapsed.as_secs_f64() * 1000.0);
}
output
};
match result {
Ok(output) => {
print!("{}", output);
ExitCode::SUCCESS
}
Err(e) => {
eprintln!("Error: {}", e);
ExitCode::FAILURE
}
}
}
fn read_stdin_json() -> Option<Arc<serde_json::Value>> {
let mut buf = String::new();
std::io::stdin().read_to_string(&mut buf).ok()?;
serde_json::from_str(&buf).ok().map(Arc::new)
}
fn handle_bench(
format: &str,
no_version: bool,
exit_code: Option<i32>,
no_color: bool,
shell: style::Shell,
stdin_data: Option<Arc<serde_json::Value>>,
) -> error::Result<String> {
let mut times = Vec::new();
for _ in 0..100 {
let start = Instant::now();
let _ = executor::execute_with_shell(
format,
no_version,
exit_code,
no_color,
shell,
stdin_data.clone(),
None,
)?;
times.push(start.elapsed());
}
times.sort();
let min = times[0];
let max = times[99];
let avg: std::time::Duration = times.iter().sum::<std::time::Duration>() / 100;
let p99 = times[98];
Ok(format!(
"100 runs: min={:.2}ms avg={:.2}ms max={:.2}ms p99={:.2}ms\n",
min.as_secs_f64() * 1000.0,
avg.as_secs_f64() * 1000.0,
max.as_secs_f64() * 1000.0,
p99.as_secs_f64() * 1000.0
))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn resolve_shell_prefers_cli_over_other_sources() {
let resolved = resolve_shell_from_sources(
Some(style::Shell::Bash),
Some(style::Shell::Zsh),
Some(style::Shell::None),
);
assert_eq!(resolved, style::Shell::Bash);
}
#[test]
fn resolve_shell_prefers_parent_process_over_env() {
let resolved =
resolve_shell_from_sources(None, Some(style::Shell::Bash), Some(style::Shell::Zsh));
assert_eq!(resolved, style::Shell::Bash);
}
#[test]
fn resolve_shell_falls_back_to_env_when_parent_missing() {
let resolved = resolve_shell_from_sources(None, None, Some(style::Shell::Zsh));
assert_eq!(resolved, style::Shell::Zsh);
}
#[test]
fn resolve_shell_defaults_to_none_when_no_source_matches() {
let resolved = resolve_shell_from_sources(None, None, None);
assert_eq!(resolved, style::Shell::None);
}
#[test]
fn parse_args_accepts_single_positional_format() {
let cli = parse_args_from(["prmt", "{path}"]).expect("parse args");
assert_eq!(cli.format.as_deref(), Some("{path}"));
}
#[test]
fn parse_args_rejects_multiple_positional_formats() {
let err = match parse_args_from(["prmt", "{path}", "{git}"]) {
Ok(_) => panic!("expected error"),
Err(err) => err,
};
assert_eq!(err.to_string(), "unexpected argument \"{git}\"");
}
#[test]
fn parse_args_rejects_positional_when_flag_format_is_present() {
let err = match parse_args_from(["prmt", "--format", "{path}", "{git}"]) {
Ok(_) => panic!("expected error"),
Err(err) => err,
};
assert_eq!(err.to_string(), "unexpected argument \"{git}\"");
}
}