1use clap::{Args, Parser, Subcommand};
2#[cfg(feature = "completion")]
3use clap::{CommandFactory, ValueEnum};
4#[cfg(feature = "completion")]
5use clap_complete::{Shell, generate};
6use signal_hook::consts::signal::{SIGINT, SIGTERM};
7use std::io::{self, IsTerminal};
8
9use crate::store::{matches_meta, MetaSelection};
10use crate::store;
11use crate::store::Meta;
12
13mod attr;
14mod ls;
15mod path;
16mod push;
17mod rm;
18
19#[derive(Parser, Debug)]
20#[command(
21 name = "stash",
22 version,
23 about = "A local store for pipeline output and ad hoc file snapshots",
24 long_about = "A local store for pipeline output and ad hoc file snapshots.\n\nWhen used without a subcommand, stash picks a mode automatically:\n - in the middle of a pipeline, it behaves like `stash tee`\n - otherwise, it behaves like `stash push`",
25 after_help = "Links:\n Docs: https://github.com/vrypan/stash/tree/main/docs\n Reference: https://github.com/vrypan/stash/blob/main/docs/reference.md\n Issues: https://github.com/vrypan/stash/issues",
26 disable_help_subcommand = true
27)]
28struct Cli {
29 #[command(flatten)]
30 push: push::PushArgs,
31
32 #[command(subcommand)]
33 command: Option<Command>,
34}
35
36#[derive(Subcommand, Debug)]
37enum Command {
38 #[command(about = "Store and return the entry key")]
39 Push(push::PushArgs),
40 #[command(about = "Store and forward to stdout")]
41 Tee(push::TeeArgs),
42 #[command(about = "Print an entry's raw data to stdout")]
43 Cat(CatArgs),
44 #[command(about = "List entries")]
45 Ls(ls::LsArgs),
46 #[command(about = "Show or update entry attributes")]
47 Attr(attr::AttrArgs),
48 #[command(about = "List attribute keys across the stash")]
49 Attrs(AttrsArgs),
50 #[command(about = "Print stash paths")]
51 Path(path::PathArgs),
52 #[command(about = "Remove entries")]
53 Rm(rm::RmArgs),
54 #[command(about = "Print the newest entry and remove it")]
55 Pop,
56 #[cfg(feature = "completion")]
57 #[command(about = "Generate shell completion scripts")]
58 Completion(CompletionArgs),
59}
60
61#[derive(Args, Debug, Clone, Default)]
62struct CatArgs {
63 #[arg(help = "Entry reference: id, n, or @n")]
64 reference: Option<String>,
65}
66
67#[derive(Args, Debug, Clone)]
68struct AttrsArgs {
69 #[arg(long, help = "Include entry count")]
70 count: bool,
71}
72
73#[cfg(feature = "completion")]
74#[derive(Args, Debug, Clone)]
75struct CompletionArgs {
76 #[arg(value_enum, help = "Shell to generate completion for")]
77 shell: CompletionShell,
78}
79
80#[cfg(feature = "completion")]
81#[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)]
82enum CompletionShell {
83 Bash,
84 Fish,
85 Zsh,
86}
87
88#[cfg(feature = "completion")]
89impl From<CompletionShell> for Shell {
90 fn from(value: CompletionShell) -> Self {
91 match value {
92 CompletionShell::Bash => Shell::Bash,
93 CompletionShell::Fish => Shell::Fish,
94 CompletionShell::Zsh => Shell::Zsh,
95 }
96 }
97}
98
99pub fn main_entry() {
100 if let Err(err) = run() {
101 if err.kind() == io::ErrorKind::BrokenPipe {
102 std::process::exit(0);
103 }
104 if let Some(partial_err) = err
105 .get_ref()
106 .and_then(|e| e.downcast_ref::<store::PartialSavedError>())
107 {
108 let code = match partial_err.signal {
109 Some(SIGINT) => 130,
110 Some(SIGTERM) => 143,
111 _ => 1,
112 };
113 std::process::exit(code);
114 }
115 eprintln!("error: {err}");
116 std::process::exit(1);
117 }
118}
119
120pub fn run() -> io::Result<()> {
121 let cli = Cli::parse();
122 match cli.command {
123 Some(Command::Push(args)) => push::push_command(args),
124 Some(Command::Tee(args)) => push::tee_command(args),
125 Some(Command::Cat(args)) => cat_command(args),
126 Some(Command::Ls(args)) => ls::ls_command(args),
127 Some(Command::Attr(args)) => attr::attr_command(args),
128 Some(Command::Attrs(args)) => attrs_command(args),
129 Some(Command::Path(args)) => path::path_command(args),
130 Some(Command::Rm(args)) => rm::rm_command(args),
131 Some(Command::Pop) => pop_command(),
132 #[cfg(feature = "completion")]
133 Some(Command::Completion(args)) => completion_command(args),
134 None => {
135 if smart_mode_uses_tee(&cli.push) {
136 push::tee_command(push::TeeArgs {
137 attr: cli.push.attr,
138 print: cli.push.print,
139 save_on_error: true,
140 })
141 } else {
142 push::push_command(cli.push)
143 }
144 }
145 }
146}
147
148fn smart_mode_uses_tee(args: &push::PushArgs) -> bool {
149 args.file.is_none() && !io::stdin().is_terminal() && !io::stdout().is_terminal()
150}
151
152fn collect_entries(sel: &MetaSelection, reverse: bool, limit: usize) -> io::Result<Vec<Meta>> {
154 let mut items = store::list()?
155 .into_iter()
156 .filter(|m| matches_meta(&m.attrs, sel))
157 .collect::<Vec<_>>();
158 if reverse {
159 items.reverse();
160 }
161 if limit > 0 && items.len() > limit {
162 items.truncate(limit);
163 }
164 Ok(items)
165}
166
167fn cat_command(args: CatArgs) -> io::Result<()> {
168 let id = if let Some(reference) = args.reference {
169 store::resolve(&reference)?
170 } else {
171 store::newest()?.id
172 };
173 let stdout = io::stdout();
174 store::cat_to_writer(&id, stdout.lock())
175}
176
177fn pop_command() -> io::Result<()> {
178 let newest = store::newest()?;
179 let stdout = io::stdout();
180 let mut out = stdout.lock();
181 store::cat_to_writer(&newest.id, &mut out)?;
182 store::remove(&newest.id)
183}
184
185fn attrs_command(args: AttrsArgs) -> io::Result<()> {
186 for (key, count) in store::all_attr_keys()? {
187 if args.count {
188 println!("{key}\t{count}");
189 } else {
190 println!("{key}");
191 }
192 }
193 Ok(())
194}
195
196#[cfg(feature = "completion")]
197fn completion_command(args: CompletionArgs) -> io::Result<()> {
198 write_completions(Shell::from(args.shell), &mut io::stdout());
199 Ok(())
200}
201
202#[cfg(feature = "completion")]
206pub fn write_completions(shell: Shell, writer: &mut dyn io::Write) {
207 let mut cmd = Cli::command();
208 let name = cmd.get_name().to_string();
209 generate(shell, &mut cmd, name, writer);
210}