Refactor + support PWD hook for zsh

src/db.rs

@@ -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;
         }

src/dir.rs

@@ -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 {

src/main.rs

@@ -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'
-"#;

src/subcommand/add.rs

@@ -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)
+            }
+        }
+    }
+}

src/subcommand/init.rs

@@ -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")
+"#;

src/subcommand/migrate.rs

@@ -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)
+    }
+}

src/subcommand/mod.rs

@@ -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;

src/subcommand/query.rs

@@ -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)
+    }
+}

src/subcommand/remove.rs

@@ -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)
+    }
+}

src/types.rs

@@ -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

src/util.rs

@@ -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))
 }