use crate::commands::{CommandContext, VirtualCommandRegistry};
use crate::{CommandBuilder, CommandResult, StdinOption};
use serde::Deserialize;
use std::collections::HashMap;
use std::future::Future;
use std::path::Path;
use std::pin::Pin;
use std::sync::Arc;
pub type HandlerFuture = Pin<Box<dyn Future<Output = CommandResult> + Send>>;
pub type PackageHandler = Arc<dyn Fn(CommandContext) -> HandlerFuture + Send + Sync>;
#[derive(Clone)]
pub struct CommandPackage {
pub name: String,
pub version: String,
pub commands: HashMap<String, PackageHandler>,
}
#[derive(Deserialize)]
struct Manifest {
name: String,
version: String,
commands: HashMap<String, Specification>,
}
#[derive(Deserialize)]
struct Specification {
file: String,
#[serde(default)]
args: Vec<String>,
}
impl CommandPackage {
pub async fn load(path: impl AsRef<Path>) -> std::io::Result<Self> {
let manifest: Manifest = serde_json::from_str(&tokio::fs::read_to_string(path).await?)?;
let commands = manifest
.commands
.into_iter()
.map(|(name, specification)| {
let command = CommandBuilder::new(specification.file).args(specification.args);
let handler: PackageHandler = Arc::new(move |context| {
let mut command = command
.clone()
.args(context.args)
.stdin(StdinOption::Content(context.stdin.unwrap_or_default()));
if let Some(cwd) = context.cwd {
command = command.cwd(cwd);
}
if let Some(env) = context.env {
for (key, value) in env {
command = command.env(key, value);
}
}
Box::pin(async move {
command.run().await.unwrap_or_else(|error| {
CommandResult::error_with_code(error.to_string(), 1)
})
})
});
(name, handler)
})
.collect();
Ok(Self {
name: manifest.name,
version: manifest.version,
commands,
})
}
}
#[derive(Default)]
pub struct CommandEnvironment {
commands: HashMap<String, PackageHandler>,
owners: HashMap<String, String>,
}
impl CommandEnvironment {
pub fn new() -> Self {
Self::default()
}
pub fn create<F, Fut>(&mut self, name: &str, handler: F) -> Result<(), String>
where
F: Fn(CommandContext) -> Fut + Send + Sync + 'static,
Fut: Future<Output = CommandResult> + Send + 'static,
{
validate(name)?;
self.owners.remove(name);
self.commands.insert(
name.into(),
Arc::new(move |context| Box::pin(handler(context))),
);
Ok(())
}
pub fn list(&self) -> Vec<String> {
let mut names = self.commands.keys().cloned().collect::<Vec<_>>();
names.sort();
names
}
pub fn get(&self, name: &str) -> Option<PackageHandler> {
self.commands.get(name).cloned().or_else(|| {
let handler = *VirtualCommandRegistry::with_builtins().get(name)?;
Some(Arc::new(handler) as PackageHandler)
})
}
pub async fn run(&self, name: &str, context: CommandContext) -> CommandResult {
match self.get(name) {
Some(handler) => handler(context).await,
None => CommandResult::error_with_code(format!("Unknown command: {name}"), 127),
}
}
pub fn extend<F>(&mut self, name: &str, middleware: F) -> Result<(), String>
where
F: Fn(CommandContext, PackageHandler) -> HandlerFuture + Send + Sync + 'static,
{
let parent = self
.get(name)
.ok_or_else(|| format!("Unknown command: {name}"))?;
self.commands.insert(
name.into(),
Arc::new(move |context| middleware(context, parent.clone())),
);
Ok(())
}
pub fn inherit(&mut self, name: &str, parent: &str, args: Vec<String>) -> Result<(), String> {
let handler = self
.get(parent)
.ok_or_else(|| format!("Unknown command: {parent}"))?;
self.create(name, move |mut context| {
let mut prefix = args.clone();
prefix.extend(context.args);
context.args = prefix;
handler(context)
})
}
pub fn compose(&mut self, name: &str, names: &[&str]) -> Result<(), String> {
if names.is_empty() {
return Err("Composition requires commands".into());
}
let handlers = names
.iter()
.map(|name| {
self.get(name)
.ok_or_else(|| format!("Unknown command: {name}"))
})
.collect::<Result<Vec<_>, _>>()?;
self.create(name, move |context| {
let handlers = handlers.clone();
Box::pin(async move {
let mut result = CommandResult::success(context.stdin.unwrap_or_default());
for handler in handlers {
let next = CommandContext {
args: context.args.clone(),
stdin: Some(result.stdout.to_string()),
cwd: context.cwd.clone(),
env: context.env.clone(),
output_tx: context.output_tx.clone(),
is_cancelled: None,
};
result = handler(next).await;
if result.code != 0 {
break;
}
}
result
})
})
}
pub fn install(&mut self, package: CommandPackage) -> Result<(), String> {
if package.commands.is_empty() {
return Err("Empty command package".into());
}
for name in package.commands.keys() {
validate(name)?;
if self.commands.contains_key(name) && self.owners.get(name) != Some(&package.name) {
return Err(format!("Command already registered: {name}"));
}
}
self.uninstall(&package.name);
for (name, handler) in package.commands {
self.owners.insert(name.clone(), package.name.clone());
self.commands.insert(name, handler);
}
Ok(())
}
pub fn uninstall(&mut self, package: &str) {
let names = self
.owners
.iter()
.filter(|(_, owner)| owner.as_str() == package)
.map(|(name, _)| name.clone())
.collect::<Vec<_>>();
for name in names {
self.owners.remove(&name);
self.commands.remove(&name);
}
}
}
fn validate(name: &str) -> Result<(), String> {
if name.is_empty()
|| !name
.chars()
.next()
.is_some_and(|c| c.is_ascii_alphabetic() || c == '_')
|| !name
.chars()
.all(|c| c.is_ascii_alphanumeric() || "_:.-".contains(c))
{
return Err(format!("Invalid command: {name}"));
}
Ok(())
}