Llewellyn/zoxide
1
0
Fork 0
mirror of https://github.com/Llewellynvdm/zoxide.git synced 2024-11-29 07:53:55 +00:00
Code Issues Projects Releases Wiki Activity

Refactor + support PWD hook for zsh

Browse Source
This commit is contained in:
Ajeet D'Souza 2020-03-13 06:19:37 +05:30
parent 4596716cc8
commit 9c8e8da71a
11 changed files with 375 additions and 263 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
src/db.rs
View File

@ -1,5 +1,5 @@
use crate::dir::Dir;
use crate::types::{Rank, Timestamp};
use crate::types::{Epoch, Rank};
use crate::util;
use anyhow::{anyhow, bail, Context, Result};
use fs2::FileExt;
@ -28,21 +28,21 @@ impl DB {
.write(true)
.create(true)
.open(&path_tmp)
.with_context(|| anyhow!("could not open temporary database"))?;
.with_context(|| anyhow!("could not open temporary database file"))?;
file_tmp
.lock_exclusive()
.with_context(|| anyhow!("could not lock temporary database"))?;
.with_context(|| anyhow!("could not lock temporary database file"))?;
let dirs = match File::open(&path) {
Ok(file) => {
let rd = BufReader::new(&file);
bincode::deserialize_from(rd)
let reader = BufReader::new(&file);
bincode::deserialize_from(reader)
.with_context(|| anyhow!("could not deserialize database"))?
}
Err(err) => match err.kind() {
io::ErrorKind::NotFound => Vec::<Dir>::new(),
_ => return Err(err).with_context(|| anyhow!("could not open database")),
_ => return Err(err).with_context(|| anyhow!("could not open database file")),
},
};
@ -59,13 +59,13 @@ impl DB {
if self.modified {
self.file_tmp
.set_len(0)
.with_context(|| "could not truncate temporary database")?;
.with_context(|| "could not truncate temporary database file")?;
let wr = BufWriter::new(&self.file_tmp);
bincode::serialize_into(wr, &self.dirs)
let writer = BufWriter::new(&self.file_tmp);
bincode::serialize_into(writer, &self.dirs)
.with_context(|| anyhow!("could not serialize database"))?;
fs::rename(&self.path_tmp, &self.path)
.with_context(|| anyhow!("could not move temporary database"))?;
.with_context(|| anyhow!("could not move temporary database file"))?;
}
Ok(())
@ -153,7 +153,7 @@ impl DB {
Ok(())
}
pub fn add<P: AsRef<Path>>(&mut self, path: P, now: Timestamp) -> Result<()> {
pub fn add<P: AsRef<Path>>(&mut self, path: P, now: Epoch) -> Result<()> {
let path_abs = path
.as_ref()
.canonicalize()
@ -165,7 +165,7 @@ impl DB {
match self.dirs.iter_mut().find(|dir| dir.path == path_str) {
None => self.dirs.push(Dir {
path: path_str.to_owned(),
path: path_str.to_string(),
last_accessed: now,
rank: 1.0,
}),
@ -191,7 +191,7 @@ impl DB {
Ok(())
}
pub fn query(&mut self, keywords: &[String], now: Timestamp) -> Option<Dir> {
pub fn query(&mut self, keywords: &[String], now: Epoch) -> Option<Dir> {
loop {
let (idx, dir) = self
.dirs
@ -209,13 +209,9 @@ impl DB {
}
}
pub fn query_all(&mut self, mut keywords: Vec<String>) -> Vec<Dir> {
pub fn query_all(&mut self, keywords: &[String]) -> Vec<Dir> {
self.remove_invalid();
for keyword in &mut keywords {
keyword.make_ascii_lowercase();
}
self.dirs
.iter()
.filter(|dir| dir.is_match(&keywords))
@ -244,6 +240,7 @@ impl DB {
fn remove_invalid(&mut self) {
let orig_len = self.dirs.len();
self.dirs.retain(Dir::is_dir);
if orig_len != self.dirs.len() {
self.modified = true;
}

12
src/dir.rs
View File

@ -1,4 +1,4 @@
use crate::types::{Rank, Timestamp};
use crate::types::{Rank, Epoch};
use serde::{Deserialize, Serialize};
use std::path::Path;
@ -6,7 +6,7 @@ use std::path::Path;
pub struct Dir {
pub path: String,
pub rank: Rank,
pub last_accessed: Timestamp,
pub last_accessed: Epoch,
}
impl Dir {
@ -40,10 +40,10 @@ impl Dir {
true
}
pub fn get_frecency(&self, now: Timestamp) -> Rank {
const HOUR: Timestamp = 60 * 60;
const DAY: Timestamp = 24 * HOUR;
const WEEK: Timestamp = 7 * DAY;
pub fn get_frecency(&self, now: Epoch) -> Rank {
const HOUR: Epoch = 60 * 60;
const DAY: Epoch = 24 * HOUR;
const WEEK: Epoch = 7 * DAY;
let duration = now - self.last_accessed;
if duration < HOUR {

245
src/main.rs
View File

@ -1,254 +1,33 @@
mod db;
mod dir;
mod subcommand;
mod types;
mod util;
use crate::util::{fzf_helper, get_current_time, get_db};
use anyhow::{anyhow, Context, Result};
use clap::arg_enum;
use std::env;
use std::path::Path;
use anyhow::Result;
use structopt::StructOpt;
// TODO: use structopt to parse env variables: <https://github.com/TeXitoi/structopt/blob/master/examples/env.rs>
arg_enum! {
#[allow(non_camel_case_types)]
#[derive(Debug)]
enum Shell {
bash,
fish,
zsh,
}
}
#[derive(Debug, StructOpt)]
#[structopt(about = "A cd command that learns your habits")]
enum Zoxide {
#[structopt(about = "Add a new directory or increment its rank")]
Add { path: Option<String> },
#[structopt(about = "Migrate from z database")]
Migrate { path: String },
#[structopt(about = "Prints shell configuration")]
Init {
#[structopt(possible_values = &Shell::variants(), case_insensitive = true)]
shell: Shell,
#[structopt(
long,
help = "Prevents zoxide from defining any aliases other than 'z'"
)]
no_define_aliases: bool,
},
#[structopt(about = "Search for a directory")]
Query {
keywords: Vec<String>,
#[structopt(short, long, help = "Opens an interactive selection menu using fzf")]
interactive: bool,
},
#[structopt(about = "Remove a directory")]
Remove { path: String },
}
fn zoxide_query(mut keywords: Vec<String>) -> Result<Option<String>> {
let now = get_current_time()?;
let mut db = get_db()?;
if let [path] = keywords.as_slice() {
if Path::new(path).is_dir() {
return Ok(Some(path.to_owned()));
}
}
for keyword in &mut keywords {
keyword.make_ascii_lowercase();
}
if let Some(dir) = db.query(&keywords, now) {
return Ok(Some(dir.path));
}
Ok(None)
}
fn zoxide_query_interactive(keywords: Vec<String>) -> Result<Option<String>> {
let now = get_current_time()?;
let dirs = get_db()?.query_all(keywords);
fzf_helper(now, dirs)
Add(subcommand::Add),
Init(subcommand::Init),
Migrate(subcommand::Migrate),
Query(subcommand::Query),
Remove(subcommand::Remove),
}
pub fn main() -> Result<()> {
let opt = Zoxide::from_args();
match opt {
Zoxide::Add { path: path_opt } => {
let mut db = get_db()?;
let now = get_current_time()?;
match path_opt {
Some(path) => db.add(path, now),
None => {
let current_dir = env::current_dir()
.with_context(|| anyhow!("unable to fetch current directory"))?;
db.add(current_dir, now)
}
}?;
}
Zoxide::Migrate { path } => {
let mut db = get_db()?;
db.migrate(path)?;
}
Zoxide::Init {
shell,
no_define_aliases,
} => {
match shell {
Shell::bash => {
println!("{}", INIT_BASH);
if !no_define_aliases {
println!("{}", INIT_BASH_ALIAS);
}
}
Shell::fish => {
println!("{}", INIT_FISH);
if !no_define_aliases {
println!("{}", INIT_FISH_ALIAS);
}
}
Shell::zsh => {
println!("{}", INIT_ZSH);
if !no_define_aliases {
println!("{}", INIT_ZSH_ALIAS);
}
}
};
}
Zoxide::Query {
keywords,
interactive,
} => {
let path_opt = if interactive {
zoxide_query_interactive(keywords)
} else {
zoxide_query(keywords)
}?;
if let Some(path) = path_opt {
println!("query: {}", path.trim());
}
}
Zoxide::Remove { path } => {
let mut db = get_db()?;
db.remove(path)?;
}
Zoxide::Add(add) => add.run()?,
Zoxide::Init(init) => init.run()?,
Zoxide::Migrate(migrate) => migrate.run()?,
Zoxide::Query(query) => query.run()?,
Zoxide::Remove(remove) => remove.run()?,
};
Ok(())
}
const INIT_BASH: &str = r#"
_zoxide_precmd() {
zoxide add
}
case "$PROMPT_COMMAND" in
*_zoxide_precmd*) ;;
*) PROMPT_COMMAND="_zoxide_precmd${PROMPT_COMMAND:+;${PROMPT_COMMAND}}" ;;
esac
z() {
if [ "${#}" -eq 0 ]; then
cd "${HOME}"
elif [ "${#}" -eq 1 ] && [ "${1}" = '-' ]; then
cd '-'
else
local result=$(zoxide query "${@}")
case "${result}" in
"query: "*) cd "${result:7}" ;;
*) [ -n "${result}" ] && echo "${result}" ;;
esac
fi
}
"#;
const INIT_BASH_ALIAS: &str = r#"
alias zi='z -i'
alias za='zoxide add'
alias zq='zoxide query'
alias zr='zoxide remove'
"#;
const INIT_FISH: &str = r#"
function _zoxide_precmd --on-event fish_prompt
zoxide add
end
function z
set -l argc (count "$argv")
if [ "$argc" -eq 0 ]
cd "$HOME"
and commandline -f repaint
else if [ "$argc" -eq 1 ]
and [ "$argv[1]" = '-' ]
cd '-'
and commandline -f repaint
else
# TODO: use string-collect from fish 3.1.0 once it has wider adoption
set -l IFS ''
set -l result (zoxide query $argv)
switch "$result"
case 'query: *'
cd (string sub -s 8 "$result")
and commandline -f repaint
case '*'
[ -n "$result" ]
and echo "$result"
end
end
end
"#;
const INIT_FISH_ALIAS: &str = r#"
abbr -a zi 'z -i'
abbr -a za 'zoxide add'
abbr -a zq 'zoxide query'
abbr -a zr 'zoxide remove'
"#;
const INIT_ZSH: &str = r#"
_zoxide_precmd() {
zoxide add
}
[[ -n "${precmd_functions[(r)_zoxide_precmd]}" ]] || {
precmd_functions+=(_zoxide_precmd)
}
z() {
if [ "${#}" -eq 0 ]; then
cd "${HOME}"
elif [ "${#}" -eq 1 ] && [ "${1}" = '-' ]; then
cd '-'
else
local result=$(zoxide query "$@")
case "$result" in
"query: "*) cd "${result:7}" ;;
*) [ -n "$result" ] && echo "$result" ;;
esac
fi
}
"#;
const INIT_ZSH_ALIAS: &str = r#"
alias zi='z -i'
alias za='zoxide add'
alias zq='zoxide query'
alias zr='zoxide remove'
"#;

26
src/subcommand/add.rs Normal file
View File

@ -0,0 +1,26 @@
use crate::util;
use anyhow::{anyhow, Context, Result};
use std::env;
use structopt::StructOpt;
#[derive(Debug, StructOpt)]
#[structopt(about = "Add a new directory or increment its rank")]
pub struct Add {
path: Option<String>,
}
impl Add {
pub fn run(&self) -> Result<()> {
let mut db = util::get_db()?;
let now = util::get_current_time()?;
match &self.path {
Some(path) => db.add(path, now),
None => {
let current_dir = env::current_dir()
.with_context(|| anyhow!("unable to fetch current directory"))?;
db.add(current_dir, now)
}
}
}
}

209
src/subcommand/init.rs Normal file
View File

@ -0,0 +1,209 @@
use anyhow::{bail, Result};
use clap::arg_enum;
use std::io::{self, Write};
use structopt::StructOpt;
#[derive(Debug, StructOpt)]
#[structopt(about = "Generates shell configuration")]
pub struct Init {
#[structopt(possible_values = &Shell::variants(), case_insensitive = true)]
shell: Shell,
#[structopt(
long,
help = "Prevents zoxide from defining any aliases other than 'z'"
)]
no_define_aliases: bool,
#[structopt(
long,
help = "Chooses event on which an entry is added to the database",
possible_values = &Hook::variants(),
default_value = "prompt",
case_insensitive = true
)]
hook: Hook,
}
impl Init {
pub fn run(&self) -> Result<()> {
let config = match self.shell {
Shell::bash => BASH_CONFIG,
Shell::fish => FISH_CONFIG,
Shell::zsh => ZSH_CONFIG,
};
let stdout = io::stdout();
let mut handle = stdout.lock();
writeln!(handle, "{}", config.z).unwrap();
if !self.no_define_aliases {
writeln!(handle, "{}", config.alias).unwrap();
}
match self.hook {
Hook::none => (),
Hook::prompt => writeln!(handle, "{}", config.hook.prompt).unwrap(),
Hook::pwd => match config.hook.pwd {
Some(pwd) => writeln!(handle, "{}", pwd).unwrap(),
None => bail!("pwd hooks are currently not supported for this shell"),
},
};
Ok(())
}
}
arg_enum! {
#[allow(non_camel_case_types)]
#[derive(Debug)]
enum Shell {
bash,
fish,
zsh,
}
}
arg_enum! {
#[allow(non_camel_case_types)]
#[derive(Debug)]
enum Hook {
none,
prompt,
pwd,
}
}
const BASH_CONFIG: ShellConfig = ShellConfig {
z: BASH_Z,
alias: BASH_ALIAS,
hook: HookConfig {
prompt: BASH_HOOK_PROMPT,
pwd: None,
},
};
const FISH_CONFIG: ShellConfig = ShellConfig {
z: FISH_Z,
alias: FISH_ALIAS,
hook: HookConfig {
prompt: FISH_HOOK_PROMPT,
pwd: None,
},
};
const ZSH_CONFIG: ShellConfig = ShellConfig {
z: ZSH_Z,
alias: ZSH_ALIAS,
hook: HookConfig {
prompt: ZSH_HOOK_PROMPT,
pwd: Some(ZSH_HOOK_PWD),
},
};
struct ShellConfig {
z: &'static str,
alias: &'static str,
hook: HookConfig,
}
struct HookConfig {
prompt: &'static str,
pwd: Option<&'static str>,
}
const BASH_Z: &str = r#"
z() {
if [ "${#}" -eq 0 ]; then
cd "${HOME}"
elif [ "${#}" -eq 1 ] && [ "${1}" = '-' ]; then
cd '-'
else
local result=$(zoxide query "${@}")
case "${result}" in
"query: "*) cd "${result:7}" ;;
*) [ -n "${result}" ] && echo "${result}" ;;
esac
fi
}
"#;
const FISH_Z: &str = r#"
function z
set -l argc (count "$argv")
if [ "$argc" -eq 0 ]
cd "$HOME"
and commandline -f repaint
else if [ "$argc" -eq 1 ]
and [ "$argv[1]" = '-' ]
cd '-'
and commandline -f repaint
else
# TODO: use string-collect from fish 3.1.0 once it has wider adoption
set -l IFS ''
set -l result (zoxide query $argv)
switch "$result"
case 'query: *'
cd (string sub -s 8 "$result")
and commandline -f repaint
case '*'
[ -n "$result" ]
and echo "$result"
end
end
end
"#;
const ZSH_Z: &str = BASH_Z;
const BASH_ALIAS: &str = r#"
alias zi='z -i'
alias za='zoxide add'
alias zq='zoxide query'
alias zr='zoxide remove'
"#;
const FISH_ALIAS: &str = r#"
abbr -a zi 'z -i'
abbr -a za 'zoxide add'
abbr -a zq 'zoxide query'
abbr -a zr 'zoxide remove'
"#;
const ZSH_ALIAS: &str = BASH_ALIAS;
const BASH_HOOK_PROMPT: &str = r#"
_zoxide_hook() {
zoxide add
}
case "$PROMPT_COMMAND" in
*_zoxide_hook*) ;;
*) PROMPT_COMMAND="_zoxide_hook${PROMPT_COMMAND:+;${PROMPT_COMMAND}}" ;;
esac
"#;
const FISH_HOOK_PROMPT: &str = r#"
function _zoxide_hook --on-event fish_prompt
zoxide add
end
"#;
const ZSH_HOOK_PROMPT: &str = r#"
_zoxide_hook() {
zoxide add
}
[[ -n "${precmd_functions[(r)_zoxide_hook]}" ]] || {
precmd_functions+=(_zoxide_hook)
}
"#;
const ZSH_HOOK_PWD: &str = r#"
_zoxide_hook() {
zoxide add
}
chpwd_functions=(${chpwd_functions[@]} "_zoxide_hook")
"#;

15
src/subcommand/migrate.rs Normal file
View File

@ -0,0 +1,15 @@
use crate::util;
use anyhow::Result;
use structopt::StructOpt;
#[derive(Debug, StructOpt)]
#[structopt(about = "Migrate from z database")]
pub struct Migrate {
path: String,
}
impl Migrate {
pub fn run(&self) -> Result<()> {
util::get_db()?.migrate(&self.path)
}
}

11
src/subcommand/mod.rs Normal file
View File

@ -0,0 +1,11 @@
mod add;
mod init;
mod migrate;
mod query;
mod remove;
pub use add::Add;
pub use init::Init;
pub use migrate::Migrate;
pub use query::Query;
pub use remove::Remove;

60
src/subcommand/query.rs Normal file
View File

@ -0,0 +1,60 @@
use crate::util;
use anyhow::Result;
use std::path::Path;
use structopt::StructOpt;
#[derive(Debug, StructOpt)]
#[structopt(about = "Search for a directory")]
pub struct Query {
keywords: Vec<String>,
#[structopt(short, long, help = "Opens an interactive selection menu using fzf")]
interactive: bool,
}
impl Query {
pub fn run(mut self) -> Result<()> {
let path_opt = if self.interactive {
self.query_interactive()
} else {
self.query()
}?;
if let Some(path) = path_opt {
println!("query: {}", path.trim());
}
Ok(())
}
fn query(&mut self) -> Result<Option<String>> {
let now = util::get_current_time()?;
let mut db = util::get_db()?;
if let [path] = self.keywords.as_slice() {
if Path::new(path).is_dir() {
return Ok(Some(path.to_string()));
}
}
for keyword in &mut self.keywords {
keyword.make_ascii_lowercase();
}
if let Some(dir) = db.query(&self.keywords, now) {
return Ok(Some(dir.path));
}
Ok(None)
}
fn query_interactive(&mut self) -> Result<Option<String>> {
let now = util::get_current_time()?;
for keyword in &mut self.keywords {
keyword.make_ascii_lowercase();
}
let dirs = util::get_db()?.query_all(&self.keywords);
util::fzf_helper(now, dirs)
}
}

15
src/subcommand/remove.rs Normal file
View File

@ -0,0 +1,15 @@
use crate::util;
use anyhow::Result;
use structopt::StructOpt;
#[derive(Debug, StructOpt)]
#[structopt(about = "Remove a directory")]
pub struct Remove {
path: String,
}
impl Remove {
pub fn run(&self) -> Result<()> {
util::get_db()?.remove(&self.path)
}
}

2
src/types.rs
View File

@ -1,3 +1,3 @@
// TODO: convert these to newtypes
pub use f64 as Rank;
pub use i64 as Timestamp; // use a signed integer so subtraction can be performed on it
pub use i64 as Epoch; // use a signed integer so subtraction can be performed on it

10
src/util.rs
View File

@ -1,6 +1,6 @@
use crate::db::DB;
use crate::dir::Dir;
use crate::types::{Rank, Timestamp};
use crate::types::{Epoch, Rank};
use anyhow::{anyhow, Context, Result};
use std::env;
use std::io::{Read, Write};
@ -42,16 +42,16 @@ pub fn get_db() -> Result<DB> {
DB::open(path)
}
pub fn get_current_time() -> Result<Timestamp> {
pub fn get_current_time() -> Result<Epoch> {
let current_time = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.with_context(|| "system clock set to invalid time")?
.as_secs();
Ok(current_time as Timestamp)
Ok(current_time as Epoch)
}
pub fn fzf_helper(now: Timestamp, mut dirs: Vec<Dir>) -> Result<Option<String>> {
pub fn fzf_helper(now: Epoch, mut dirs: Vec<Dir>) -> Result<Option<String>> {
let fzf = Command::new("fzf")
.arg("-n2..")
.stdin(Stdio::piped())
@ -92,5 +92,5 @@ pub fn fzf_helper(now: Timestamp, mut dirs: Vec<Dir>) -> Result<Option<String>>
.read_to_string(&mut output)
.with_context(|| anyhow!("could not read from fzf stdout"))?;
Ok(output.get(12..).map(str::to_owned))
Ok(output.get(12..).map(str::to_string))
}