mod shells;
use std::ffi::OsString;
use std::io::Write as _;
pub use shells::*;
pub struct CompleteEnv<'s, F> {
factory: F,
var: &'static str,
bin: Option<String>,
completer: Option<String>,
shells: Shells<'s>,
}
impl<'s, F: Fn() -> clap::Command> CompleteEnv<'s, F> {
pub fn with_factory(factory: F) -> Self {
Self {
factory,
var: "COMPLETE",
bin: None,
completer: None,
shells: Shells::builtins(),
}
}
pub fn var(mut self, var: &'static str) -> Self {
self.var = var;
self
}
pub fn bin(mut self, bin: impl Into<String>) -> Self {
self.bin = Some(bin.into());
self
}
pub fn completer(mut self, completer: impl Into<String>) -> Self {
self.completer = Some(completer.into());
self
}
pub fn shells(mut self, shells: Shells<'s>) -> Self {
self.shells = shells;
self
}
}
impl<'s, F: Fn() -> clap::Command> CompleteEnv<'s, F> {
pub fn complete(self) {
let args = std::env::args_os();
let current_dir = std::env::current_dir().ok();
if self
.try_complete(args, current_dir.as_deref())
.unwrap_or_else(|e| e.exit())
{
std::process::exit(0)
}
}
pub fn try_complete(
self,
args: impl IntoIterator<Item = impl Into<OsString>>,
current_dir: Option<&std::path::Path>,
) -> clap::error::Result<bool> {
self.try_complete_(args.into_iter().map(|a| a.into()).collect(), current_dir)
}
fn try_complete_(
self,
mut args: Vec<OsString>,
current_dir: Option<&std::path::Path>,
) -> clap::error::Result<bool> {
let Some(name) = std::env::var_os(self.var) else {
return Ok(false);
};
if name.is_empty() || name == "0" {
return Ok(false);
}
unsafe {
std::env::remove_var(self.var);
}
let shell = self.shell(std::path::Path::new(&name))?;
let mut cmd = (self.factory)();
cmd.build();
let completer = args.remove(0);
let escape_index = args
.iter()
.position(|a| *a == "--")
.map(|i| i + 1)
.unwrap_or(args.len());
args.drain(0..escape_index);
if args.is_empty() {
let mut buf = Vec::new();
self.write_registration(&cmd, current_dir, shell, completer, &mut buf)?;
std::io::stdout().write_all(&buf)?;
} else {
let mut buf = Vec::new();
shell.write_complete(&mut cmd, args, current_dir, &mut buf)?;
std::io::stdout().write_all(&buf)?;
}
Ok(true)
}
fn shell(&self, name: &std::path::Path) -> Result<&dyn EnvCompleter, std::io::Error> {
let name = name.file_stem().unwrap_or(name.as_os_str());
let name = name.to_string_lossy();
let shell = self.shells.completer(&name).ok_or_else(|| {
let shells =
self.shells
.names()
.enumerate()
.fold(String::new(), |mut seed, (i, name)| {
use std::fmt::Write as _;
let prefix = if i == 0 { "" } else { ", " };
let _ = write!(&mut seed, "{prefix}`{name}`");
seed
});
std::io::Error::other(format!("unknown shell `{name}`, expected one of {shells}"))
})?;
Ok(shell)
}
fn write_registration(
&self,
cmd: &clap::Command,
current_dir: Option<&std::path::Path>,
shell: &dyn EnvCompleter,
completer: OsString,
buf: &mut dyn std::io::Write,
) -> Result<(), std::io::Error> {
let name = cmd.get_name();
let bin = self
.bin
.as_deref()
.or_else(|| cmd.get_bin_name())
.unwrap_or_else(|| cmd.get_name());
let completer = if let Some(completer) = self.completer.as_deref() {
completer.to_owned()
} else {
let mut completer = std::path::PathBuf::from(completer);
if let Some(current_dir) = current_dir {
if 1 < completer.components().count() {
completer = current_dir.join(completer);
}
}
completer.to_string_lossy().into_owned()
};
shell.write_registration(self.var, name, bin, &completer, buf)?;
Ok(())
}
}
pub struct Shells<'s>(pub &'s [&'s dyn EnvCompleter]);
impl<'s> Shells<'s> {
pub const fn builtins() -> Self {
Self(&[&Bash, &Elvish, &Fish, &Powershell, &Zsh])
}
pub fn completer(&self, name: &str) -> Option<&dyn EnvCompleter> {
self.0.iter().copied().find(|c| c.is(name))
}
pub fn names(&self) -> impl Iterator<Item = &'static str> + 's {
self.0.iter().map(|c| c.name())
}
pub fn iter(&self) -> impl Iterator<Item = &dyn EnvCompleter> {
self.0.iter().copied()
}
}
pub trait EnvCompleter {
fn name(&self) -> &'static str;
fn is(&self, name: &str) -> bool;
fn write_registration(
&self,
var: &str,
name: &str,
bin: &str,
completer: &str,
buf: &mut dyn std::io::Write,
) -> Result<(), std::io::Error>;
fn write_complete(
&self,
cmd: &mut clap::Command,
args: Vec<OsString>,
current_dir: Option<&std::path::Path>,
buf: &mut dyn std::io::Write,
) -> Result<(), std::io::Error>;
}