Repository: j178/prefligit (language: python)

Issue:
Stash stage dirty files and restore



Discussion / hints:

Fix patch:
diff --git a/src/cleanup.rs b/src/cleanup.rs
new file mode 100644
index 0000000..03db74b
--- /dev/null
+++ b/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));
+}
diff --git a/src/cli/mod.rs b/src/cli/mod.rs
index aedff20..1efae69 100644
--- a/src/cli/mod.rs
+++ b/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>,
diff --git a/src/cli/run.rs b/src/cli/run.rs
index 53b1515..e84d67f 100644
--- a/src/cli/run.rs
+++ b/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,7 +288,7 @@ 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);
     }
@@ -293,7 +296,7 @@ async fn all_filenames(
     // 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!(
diff --git a/src/git.rs b/src/git.rs
index 73a1d04..db471c6 100644
--- a/src/git.rs
+++ b/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")
diff --git a/src/main.rs b/src/main.rs
index a108252..28e6a9f 100644
--- a/src/main.rs
+++ b/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
diff --git a/src/run.rs b/src/run.rs
index 3006879..3a70c0c 100644
--- a/src/run.rs
+++ b/src/run.rs
@@ -3,10 +3,11 @@ use std::collections::HashMap;
 use std::fmt::Write as _;
 use std::future::Future;
 use std::io::Write as _;
-use std::path::Path;
-use std::sync::Arc;
+use std::path::{Path, PathBuf};
+use std::process::Command;
+use std::sync::{Arc, Mutex};
 
-use anstream::ColorChoice;
+use anstream::{eprintln, ColorChoice};
 use anyhow::Result;
 use fancy_regex::{self as regex, Regex};
 use owo_colors::{OwoColorize, Style};
@@ -17,12 +18,15 @@ use tokio::task::JoinSet;
 use tracing::{error, trace};
 use unicode_width::UnicodeWidthStr;
 
+use crate::cleanup::add_cleanup;
 use crate::cli::ExitStatus;
-use crate::git::{get_diff, GIT};
+use crate::fs::Simplified;
+use crate::git;
+use crate::git::{get_diff, git_cmd, GIT};
 use crate::hook::Hook;
 use crate::identify::tags_from_path;
 use crate::printer::Printer;
-use crate::process::Cmd;
+use crate::store::Store;
 
 const SKIPPED: &str = "Skipped";
 const NO_FILES: &str = "(no files to check)";
@@ -164,7 +168,7 @@ pub async fn run_hooks(
             ColorChoice::Always | ColorChoice::AlwaysAnsi => "--color=always",
             ColorChoice::Never => "--color=never",
         };
-        Cmd::new(GIT.as_ref()?, "run git diff")
+        git_cmd("git diff")?
             .arg("--no-pager")
             .arg("diff")
             .arg("--no-ext-diff")
@@ -385,15 +389,11 @@ fn partitions<'a>(
     partitions
 }
 
-pub async fn run_by_batch<T, F, Fut>(
-    hook: &Hook,
-    filenames: &[&String],
-    run: F,
-) -> anyhow::Result<Vec<T>>
+pub async fn run_by_batch<T, F, Fut>(hook: &Hook, filenames: &[&String], run: F) -> Result<Vec<T>>
 where
     F: Fn(Vec<String>) -> Fut,
     F: Clone + Send + Sync + 'static,
-    Fut: Future<Output = anyhow::Result<T>> + Send + 'static,
+    Fut: Future<Output = Result<T>> + Send + 'static,
     T: Send + 'static,
 {
     let mut concurrency = target_concurrency(hook.require_serial);
@@ -438,3 +438,239 @@ where
 
     Ok(results)
 }
+
+static RESTORE_WORKTREE: Mutex<Option<WorkTreeKeeper>> = Mutex::new(None);
+
+struct IntentToAddKeeper(Vec<PathBuf>);
+struct WorkingTreeKeeper(Option<PathBuf>);
+
+impl IntentToAddKeeper {
+    async fn clean() -> Result<Self> {
+        let files = git::intent_to_add_files().await?;
+        if files.is_empty() {
+            return Ok(Self(vec![]));
+        }
+
+        // TODO: xargs
+        git_cmd("git rm")?
+            .arg("rm")
+            .arg("--cached")
+            .arg("--")
+            .args(&files)
+            .check(true)
+            .stdout(std::process::Stdio::null())
+            .stderr(std::process::Stdio::null())
+            .status()
+            .await?;
+
+        Ok(Self(files.into_iter().map(PathBuf::from).collect()))
+    }
+
+    fn restore(&self) -> Result<()> {
+        // Restore the intent-to-add changes.
+        if !self.0.is_empty() {
+            Command::new(GIT.as_ref()?)
+                .arg("add")
+                .arg("--intent-to-add")
+                .arg("--")
+                // TODO: xargs
+                .args(&self.0)
+                .stdout(std::process::Stdio::null())
+                .stderr(std::process::Stdio::null())
+                .status()?;
+        }
+        Ok(())
+    }
+}
+
+impl Drop for IntentToAddKeeper {
+    fn drop(&mut self) {
+        if let Err(err) = self.restore() {
+            eprintln!(
+                "{}",
+                format!("Failed to restore intent-to-add changes: {err}").red()
+            );
+        }
+    }
+}
+
+impl WorkingTreeKeeper {
+    async fn clean(patch_dir: &Path) -> Result<Self> {
+        let tree = git::write_tree().await?;
+
+        let mut cmd = git_cmd("git diff-index")?;
+        let output = cmd
+            .arg("diff-index")
+            .arg("--ignore-submodules")
+            .arg("--binary")
+            .arg("--exit-code")
+            .arg("--no-color")
+            .arg("--no-ext-diff")
+            .arg(tree)
+            .arg("--")
+            .check(false)
+            .output()
+            .await?;
+
+        if output.status.success() {
+            trace!("No non-staged changes detected");
+            // No non-staged changes
+            Ok(Self(None))
+        } else if output.status.code() == Some(1) {
+            if output.stdout.trim_ascii().is_empty() {
+                trace!("diff-index status code 1 with empty stdout");
+                // probably git auto crlf behavior quirks
+                Ok(Self(None))
+            } else {
+                let now = std::time::SystemTime::now();
+                let pid = std::process::id();
+                let patch_name = format!(
+                    "{}-{}.patch",
+                    now.duration_since(std::time::UNIX_EPOCH)?.as_millis(),
+                    pid
+                );
+                let patch_path = patch_dir.join(&patch_name);
+
+                eprintln!(
+                    "{}",
+                    format!(
+                        "Non-staged changes detected, saving to `{}`",
+                        patch_path.user_display()
+                    )
+                    .yellow()
+                );
+                fs_err::create_dir_all(patch_dir)?;
+                fs_err::write(&patch_path, output.stdout)?;
+
+                // Clean the working tree
+                Self::checkout_working_tree()?;
+
+                Ok(Self(Some(patch_path)))
+            }
+        } else {
+            Err(cmd.check_status(output.status).unwrap_err().into())
+        }
+    }
+
+    fn checkout_working_tree() -> Result<()> {
+        let status = Command::new(GIT.as_ref()?)
+            .arg("-c")
+            .arg("submodule.recurse=0")
+            .arg("checkout")
+            .arg("--")
+            .arg(".")
+            // prevent recursive post-checkout hooks
+            .env("_PRE_COMMIT_SKIP_POST_CHECKOUT", "1")
+            .stdout(std::process::Stdio::null())
+            .stderr(std::process::Stdio::null())
+            .status()?;
+        if status.success() {
+            Ok(())
+        } else {
+            Err(anyhow::anyhow!("Failed to checkout working tree"))
+        }
+    }
+
+    fn git_apply(patch: &Path) -> Result<()> {
+        let status = Command::new(GIT.as_ref()?)
+            .arg("apply")
+            .arg("--whitespace=nowarn")
+            .arg(patch)
+            .stdout(std::process::Stdio::null())
+            .stderr(std::process::Stdio::null())
+            .status()?;
+        if status.success() {
+            Ok(())
+        } else {
+            Err(anyhow::anyhow!("Failed to apply the patch"))
+        }
+    }
+
+    fn restore(&self) -> Result<()> {
+        let Some(patch) = self.0.as_ref() else {
+            return Ok(());
+        };
+
+        // Try to apply the patch
+        if Self::git_apply(patch).is_err() {
+            error!("Failed to apply the patch, rolling back changes");
+            eprintln!(
+                "{}",
+                "Failed to apply the patch, rolling back changes".red()
+            );
+
+            Self::checkout_working_tree()?;
+            Self::git_apply(patch)?;
+        };
+
+        eprintln!(
+            "{}",
+            format!(
+                "\nRestored working tree changes from `{}`",
+                patch.user_display()
+            )
+            .yellow()
+        );
+
+        Ok(())
+    }
+}
+
+impl Drop for WorkingTreeKeeper {
+    fn drop(&mut self) {
+        if let Err(err) = self.restore() {
+            eprintln!(
+                "{}",
+                format!("Failed to restore working tree changes: {err}").red()
+            );
+        }
+    }
+}
+
+/// Clean Git intent-to-add files and working tree changes, and restore them when dropped.
+pub struct WorkTreeKeeper {
+    intent_to_add: Option<IntentToAddKeeper>,
+    working_tree: Option<WorkingTreeKeeper>,
+}
+
+#[derive(Default)]
+pub struct RestoreGuard {
+    _guard: (),
+}
+
+impl Drop for RestoreGuard {
+    fn drop(&mut self) {
+        if let Some(mut keeper) = RESTORE_WORKTREE.lock().unwrap().take() {
+            keeper.restore();
+        }
+    }
+}
+
+impl WorkTreeKeeper {
+    /// Clear intent-to-add changes from the index and clear the non-staged changes from the working directory.
+    /// Restore them when the instance is dropped.
+    pub async fn clean(store: &Store) -> Result<RestoreGuard> {
+        let cleaner = Self {
+            intent_to_add: Some(IntentToAddKeeper::clean().await?),
+            working_tree: Some(WorkingTreeKeeper::clean(store.path()).await?),
+        };
+
+        // Set to the global for the cleanup hook.
+        *RESTORE_WORKTREE.lock().unwrap() = Some(cleaner);
+
+        // Make sure restoration when ctrl-c is pressed.
+        add_cleanup(|| {
+            if let Some(guard) = &mut *RESTORE_WORKTREE.lock().unwrap() {
+                guard.restore();
+            }
+        });
+
+        Ok(RestoreGuard::default())
+    }
+
+    /// Restore the intent-to-add changes and non-staged changes.
+    fn restore(&mut self) {
+        self.intent_to_add.take();
+        self.working_tree.take();
+    }
+}


Test patch:
diff --git a/tests/common/mod.rs b/tests/common/mod.rs
index 550f3c4..a0cccb3 100644
--- a/tests/common/mod.rs
+++ b/tests/common/mod.rs
@@ -179,26 +179,6 @@ impl TestContext {
 
     /// Initialize a sample project for pre-commit.
     pub fn init_project(&self) {
-        // TODO: Clone some repositories used in the tests.
-        // if !self.cache_dir.join("pre-commit-hooks") {
-        //     Command::new("git")
-        //         .arg("clone")
-        //         .arg("https://github.com/pre-commit/pre-commit-hooks")
-        //         .arg("--depth=1")
-        //         .arg("--branch=v5.0.0")
-        //         .current_dir(&self.cache_dir)
-        //         .assert()
-        //         .success();
-        // }
-        //
-        // fs_extra::dir::copy(
-        //     self.cache_dir.join("pre-commit-hooks"),
-        //     self.home_dir.join("pre-commit-hooks"),
-        //     &fs_extra::dir::CopyOptions::new(),
-        // )?;
-
-        // Write some common files.
-
         Command::new("git")
             .arg("init")
             .current_dir(&self.temp_dir)
@@ -260,6 +240,8 @@ pub const INSTA_FILTERS: &[(&str, &str)] = &[
         r"Caused by: .* \(os error 2\)",
         "Caused by: No such file or directory (os error 2)",
     ),
+    // Time seconds
+    (r"(\d+\.)?\d+s", "[TIME]"),
 ];
 
 #[allow(unused_macros)]
diff --git a/tests/files/Dockerfile b/tests/files/Dockerfile
deleted file mode 100644
index e69de29..0000000
diff --git a/tests/run.rs b/tests/run.rs
index 52fa80d..9ca19d5 100644
--- a/tests/run.rs
+++ b/tests/run.rs
@@ -1,5 +1,6 @@
 use anyhow::Result;
 use assert_fs::prelude::*;
+use insta::assert_snapshot;
 
 use crate::common::{cmd_snapshot, TestContext};
 
@@ -52,6 +53,8 @@ fn run_basic() -> Result<()> {
     ----- stderr -----
     "#);
 
+    context.git_add(".");
+
     cmd_snapshot!(context.filters(), context.run().arg("trailing-whitespace"), @r#"
     success: true
     exit_code: 0
@@ -61,6 +64,8 @@ fn run_basic() -> Result<()> {
     ----- stderr -----
     "#);
 
+    context.git_add(".");
+
     cmd_snapshot!(context.filters(), context.run().arg("typos").arg("--hook-stage").arg("pre-push"), @r#"
     success: false
     exit_code: 1
@@ -600,3 +605,112 @@ fn pass_env_vars() {
     let env = context.read("env.txt");
     assert_eq!(env, "1\n");
 }
+
+#[test]
+fn staged_files_only() -> Result<()> {
+    let context = TestContext::new();
+    context.init_project();
+    context.write_pre_commit_config(indoc::indoc! {r#"
+        repos:
+          - repo: local
+            hooks:
+              - id: trailing-whitespace
+                name: trailing-whitespace
+                language: system
+                entry: python3 -c 'print(open("file.txt", "rt").read())'
+                verbose: true
+                types: [text]
+   "#});
+
+    context
+        .workdir()
+        .child("file.txt")
+        .write_str("Hello, world!")?;
+    context.git_add(".");
+
+    // Non-staged files should be stashed and restored.
+    context
+        .workdir()
+        .child("file.txt")
+        .write_str("Hello world again!")?;
+
+    let filters: Vec<_> = context
+        .filters()
+        .into_iter()
+        .chain([(r"/\d+-\d+.patch", "/[TIME]-[PID].patch")])
+        .collect();
+
+    cmd_snapshot!(filters, context.run(), @r#"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    trailing-whitespace......................................................Passed
+    - hook id: trailing-whitespace
+    - duration: [TIME]
+      Hello, world!
+
+    ----- stderr -----
+    Non-staged changes detected, saving to `[HOME]/[TIME]-[PID].patch`
+
+    Restored working tree changes from `[HOME]/[TIME]-[PID].patch`
+    "#);
+
+    let content = context.read("file.txt");
+    assert_snapshot!(content, @"Hello world again!");
+
+    Ok(())
+}
+
+#[cfg(unix)]
+#[test]
+fn restore_on_interrupt() -> Result<()> {
+    let context = TestContext::new();
+    context.init_project();
+    // The hook will sleep for 3 seconds.
+    context.write_pre_commit_config(indoc::indoc! {r#"
+        repos:
+          - repo: local
+            hooks:
+              - id: trailing-whitespace
+                name: trailing-whitespace
+                language: system
+                entry: python3 -c 'import time; open("out.txt", "wt").write(open("file.txt", "rt").read()); time.sleep(10)'
+                verbose: true
+                types: [text]
+   "#});
+
+    context
+        .workdir()
+        .child("file.txt")
+        .write_str("Hello, world!")?;
+    context.git_add(".");
+
+    // Non-staged files should be stashed and restored.
+    context
+        .workdir()
+        .child("file.txt")
+        .write_str("Hello world again!")?;
+
+    let mut child = context.run().spawn()?;
+    let child_id = child.id();
+
+    // Send an interrupt signal to the process.
+    let handle = std::thread::spawn(move || {
+        std::thread::sleep(std::time::Duration::from_secs(1));
+        #[allow(clippy::cast_possible_wrap)]
+        unsafe {
+            libc::kill(child_id as i32, libc::SIGINT)
+        };
+    });
+
+    handle.join().unwrap();
+    child.wait()?;
+
+    let content = context.read("out.txt");
+    assert_snapshot!(content, @"Hello, world!");
+
+    let content = context.read("file.txt");
+    assert_snapshot!(content, @"Hello world again!");
+
+    Ok(())
+}