Stash working tree before running hooks (#96)

* Add test

* Fix test

* add signal test

* Fix test

* Fix tests

* Fix tests

src/cleanup.rs

@@ -0,0 +1,17 @@
+use std::sync::Mutex;
+
+static CLEANUP_HOOKS: Mutex<Vec<Box<dyn Fn() + Send>>> = Mutex::new(Vec::new());
+
+/// Run all cleanup functions.
+pub fn cleanup() {
+    let mut cleanup = CLEANUP_HOOKS.lock().unwrap();
+    for f in cleanup.drain(..) {
+        f();
+    }
+}
+
+/// Add a cleanup function to be run when the program is interrupted.
+pub fn add_cleanup<F: Fn() + Send + 'static>(f: F) {
+    let mut cleanup = CLEANUP_HOOKS.lock().unwrap();
+    cleanup.push(Box::new(f));
+}

src/cli/mod.rs

@@ -219,7 +219,7 @@ pub(crate) struct RunExtraArgs {
     pub(crate) remote_branch: Option<String>,
     #[arg(long, hide = true)]
     pub(crate) local_branch: Option<String>,
-    #[arg(long, hide = true)]
+    #[arg(long, hide = true, required_if_eq("hook_stage", "pre-rebase"))]
     pub(crate) pre_rebase_upstream: Option<String>,
     #[arg(long, hide = true)]
     pub(crate) pre_rebase_branch: Option<String>,

src/cli/run.rs

@@ -13,11 +13,10 @@ use tracing::{debug, trace};
 use crate::cli::{ExitStatus, RunExtraArgs};
 use crate::config::Stage;
 use crate::fs::{normalize_path, Simplified};
-use crate::git::{get_all_files, get_changed_files, get_staged_files, has_unmerged_paths, GIT};
+use crate::git;
 use crate::hook::{Hook, Project};
 use crate::printer::Printer;
-use crate::process::Cmd;
-use crate::run::{run_hooks, FilenameFilter};
+use crate::run::{run_hooks, FilenameFilter, WorkTreeKeeper};
 use crate::store::Store;
 
 #[allow(clippy::too_many_arguments)]
@@ -34,10 +33,17 @@ pub(crate) async fn run(
     verbose: bool,
     printer: Printer,
 ) -> Result<ExitStatus> {
+    // Prevent recursive post-checkout hooks.
+    if matches!(hook_stage, Some(Stage::PostCheckout))
+        && std::env::var_os("_PRE_COMMIT_SKIP_POST_CHECKOUT").is_some()
+    {
+        return Ok(ExitStatus::Success);
+    }
+
     let should_stash = !all_files && files.is_empty();
 
     // Check if we have unresolved merge conflict files and fail fast.
-    if should_stash && has_unmerged_paths().await? {
+    if should_stash && git::has_unmerged_paths().await? {
         writeln!(
             printer.stderr(),
             "You have unmerged paths. Resolve them before running pre-commit."
@@ -55,21 +61,12 @@ pub(crate) async fn run(
         return Ok(ExitStatus::Failure);
     }
 
-    // Prevent recursive post-checkout hooks.
-    if matches!(hook_stage, Some(Stage::PostCheckout))
-        && std::env::var_os("_PRE_COMMIT_SKIP_POST_CHECKOUT").is_some()
-    {
-        return Ok(ExitStatus::Success);
-    }
-
+    // Set env vars for hooks.
     let env_vars = fill_envs(from_ref.as_ref(), to_ref.as_ref(), &extra_args);
 
     let mut project = Project::new(config_file)?;
     let store = Store::from_settings()?.init()?;
 
-    // TODO: fill env vars
-    // TODO: impl staged_files_only
-
     let lock = store.lock_async().await?;
     let hooks = project.init_hooks(&store, printer).await?;
 
@@ -123,6 +120,12 @@ pub(crate) async fn run(
     install_hooks(&to_run, printer).await?;
     drop(lock);
 
+    // Clear any unstaged changes from the git working directory.
+    let mut _guard = None;
+    if should_stash {
+        _guard = Some(WorkTreeKeeper::clean(&store).await?);
+    }
+
     let mut filenames = all_filenames(
         hook_stage,
         from_ref,
@@ -167,7 +170,7 @@ pub(crate) async fn run(
 }
 
 async fn config_not_staged(config: &Path) -> Result<bool> {
-    let status = Cmd::new(GIT.as_ref()?, "git diff")
+    let status = git::git_cmd("git diff")?
         .arg("diff")
         .arg("--quiet") // Implies --exit-code
         .arg("--no-ext-diff")
@@ -266,7 +269,7 @@ async fn all_filenames(
             .to_string()]);
     }
     if let (Some(from_ref), Some(to_ref)) = (from_ref, to_ref) {
-        let files = get_changed_files(&from_ref, &to_ref).await?;
+        let files = git::get_changed_files(&from_ref, &to_ref).await?;
         debug!(
             "Files changed between {} and {}: {}",
             from_ref,
@@ -285,15 +288,15 @@ async fn all_filenames(
         return Ok(files);
     }
     if all_files {
-        let files = get_all_files().await?;
+        let files = git::get_all_files().await?;
         debug!("All files in the repo: {}", files.len());
         return Ok(files);
     }
     // TODO: implement merge conflict
     // if is_in_merge_conflict() {
     //     return get_conflicted_files();
     // }
-    let files = get_staged_files().await?;
+    let files = git::get_staged_files().await?;
     debug!("Staged files: {}", files.len());
     Ok(files)
 }
@@ -304,7 +307,7 @@ async fn install_hook(hook: &Hook, env_dir: PathBuf, printer: Printer) -> Result
         "Installing environment for {}",
         hook.repo(),
     )?;
-    debug!("Install environment for {} to {}", hook, env_dir.display());
+    debug!(%hook, target = %env_dir.display(), "Install environment");
 
     if env_dir.try_exists()? {
         debug!(

src/git.rs

@@ -40,7 +40,7 @@ static GIT_ENV: LazyLock<Vec<(String, String)>> = LazyLock::new(|| {
         .collect()
 });
 
-fn git_cmd(summary: &str) -> Result<Cmd, Error> {
+pub fn git_cmd(summary: &str) -> Result<Cmd, Error> {
     let mut cmd = Cmd::new(GIT.as_ref().map_err(|&e| Error::GitNotFound(e))?, summary);
     cmd.arg("-c").arg("core.useBuiltinFSMonitor=false");
     cmd.envs(GIT_ENV.iter().cloned());
@@ -60,7 +60,20 @@ fn zsplit(s: &[u8]) -> Vec<String> {
     }
 }
 
-// TODO: improve error display
+pub async fn intent_to_add_files() -> Result<Vec<String>, Error> {
+    let output = git_cmd("get intent to add files")?
+        .arg("diff")
+        .arg("--no-ext-diff")
+        .arg("--ignore-submodules")
+        .arg("--diff-filter=A")
+        .arg("--name-only")
+        .arg("-z")
+        .check(true)
+        .output()
+        .await?;
+    Ok(zsplit(&output.stdout))
+}
+
 pub async fn get_changed_files(old: &str, new: &str) -> Result<Vec<String>, Error> {
     let output = git_cmd("get changed files")?
         .arg("diff")

src/main.rs

@@ -11,10 +11,12 @@ use tracing::{debug, error};
 use tracing_subscriber::filter::Directive;
 use tracing_subscriber::EnvFilter;
 
+use crate::cleanup::cleanup;
 use crate::cli::{Cli, Command, ExitStatus, SelfCommand, SelfNamespace, SelfUpdateArgs};
 use crate::git::get_root;
 use crate::printer::Printer;
 
+mod cleanup;
 mod cli;
 mod config;
 mod fs;
@@ -248,6 +250,8 @@ async fn run(mut cli: Cli) -> Result<ExitStatus> {
 
 fn main() -> ExitCode {
     ctrlc::set_handler(move || {
+        cleanup();
+
         #[allow(clippy::exit, clippy::cast_possible_wrap)]
         std::process::exit(if cfg!(windows) {
             0xC000_013A_u32 as i32