From 0d19b4b7fd85a695efe49ae38a7afc9e85da86eb Mon Sep 17 00:00:00 2001
From: shahin <shahin@users.noreply.github.com>
Date: Thu, 18 Aug 2022 10:57:09 -0700
Subject: [PATCH] Add automatic addition of changes from pre-commit to the
 commit

---
 pre_commit/clientlib.py         |  1 +
 pre_commit/commands/run.py      | 21 ++++++++++++++++++---
 pre_commit/git.py               |  5 +++++
 pre_commit/hook.py              |  1 +
 pre_commit/staged_files_only.py | 28 +++++++++++++++++++++++++++-
 5 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/pre_commit/clientlib.py b/pre_commit/clientlib.py
index 9b53e810..f8d2658e 100644
--- a/pre_commit/clientlib.py
+++ b/pre_commit/clientlib.py
@@ -81,6 +81,7 @@ MANIFEST_HOOK_DICT = cfgv.Map(
     cfgv.Optional('require_serial', cfgv.check_bool, False),
     cfgv.Optional('stages', cfgv.check_array(cfgv.check_one_of(C.STAGES)), []),
     cfgv.Optional('verbose', cfgv.check_bool, False),
+    cfgv.Optional('commit_changes', cfgv.check_bool, False),
 )
 MANIFEST_SCHEMA = cfgv.Array(MANIFEST_HOOK_DICT)
 
diff --git a/pre_commit/commands/run.py b/pre_commit/commands/run.py
index 37f78f74..7937febc 100644
--- a/pre_commit/commands/run.py
+++ b/pre_commit/commands/run.py
@@ -164,6 +164,7 @@ def _run_single_hook(
         retcode = 0
         diff_after = diff_before
         files_modified = False
+        hook_failed = False
         out = b''
     elif not filenames and not hook.always_run:
         output.write(
@@ -180,6 +181,7 @@ def _run_single_hook(
         retcode = 0
         diff_after = diff_before
         files_modified = False
+        hook_failed = False
         out = b''
     else:
         # print hook and dots first in case the hook takes a while to run
@@ -193,10 +195,23 @@ def _run_single_hook(
         duration = round(time.time() - time_before, 2) or 0
         diff_after = _get_diff()
 
-        # if the hook makes changes, fail the commit
+        # if the hook makes changes, fail the commit or add the changes
+        # automatically
         files_modified = diff_before != diff_after
+        # We can't commit changes if another hook has made uncommitted
+        # modifications.
+        # In that case, fail the hook instead.
+        can_commit_changes = hook.commit_changes and not diff_before
+        hook_failed = bool(retcode) or (
+            files_modified and not can_commit_changes
+        )
 
-        if retcode or files_modified:
+        if files_modified and can_commit_changes:
+            git.update_changes()
+            diff_after = _get_diff()
+            assert not diff_after, 'Found unstaged diff'
+
+        if hook_failed:
             print_color = color.RED
             status = 'Failed'
         else:
@@ -223,7 +238,7 @@ def _run_single_hook(
             output.write_line_b(out.strip(), logfile_name=hook.log_file)
             output.write_line()
 
-    return files_modified or bool(retcode), diff_after
+    return hook_failed, diff_after
 
 
 def _compute_cols(hooks: Sequence[Hook]) -> int:
diff --git a/pre_commit/git.py b/pre_commit/git.py
index 35392b34..761a9d8c 100644
--- a/pre_commit/git.py
+++ b/pre_commit/git.py
@@ -215,6 +215,11 @@ def commit(repo: str = '.') -> None:
     cmd_output_b(*cmd, cwd=repo, env=env)
 
 
+def update_changes(repo: str = '.') -> None:
+    cmd = ('git', 'add', '--update')
+    cmd_output_b(*cmd, cwd=repo)
+
+
 def git_path(name: str, repo: str = '.') -> str:
     _, out, _ = cmd_output('git', 'rev-parse', '--git-path', name, cwd=repo)
     return os.path.join(repo, out.strip())
diff --git a/pre_commit/hook.py b/pre_commit/hook.py
index 202abb35..91a37073 100644
--- a/pre_commit/hook.py
+++ b/pre_commit/hook.py
@@ -36,6 +36,7 @@ class Hook(NamedTuple):
     require_serial: bool
     stages: Sequence[str]
     verbose: bool
+    commit_changes: bool
 
     @property
     def cmd(self) -> tuple[str, ...]:
diff --git a/pre_commit/staged_files_only.py b/pre_commit/staged_files_only.py
index 83d8a03e..807f48a9 100644
--- a/pre_commit/staged_files_only.py
+++ b/pre_commit/staged_files_only.py
@@ -82,9 +82,35 @@ def _unstaged_changes_cleared(patch_dir: str) -> Generator[None, None, None]:
                 # We failed to apply the patch, presumably due to fixes made
                 # by hooks.
                 # Roll back the changes made by hooks.
-                cmd_output_b(*_CHECKOUT_CMD, env=no_checkout_env)
+                cmd_output_b(
+                    'git',
+                    'restore',
+                    '--source',
+                    tree,
+                    '--worktree',
+                    '.',
+                    env=no_checkout_env,
+                )
                 _git_apply(patch_filename)
 
+                # Save the current satte of the index which may include
+                # partially staged files, a new commit, etc.
+                new_tree = cmd_output('git', 'write-tree')[1].strip()
+                # Restore worktree to the patch base (which updates the index
+                # as a side-effect).
+                cmd_output_b(
+                    'git',
+                    'checkout',
+                    tree,
+                    '--',
+                    '.',
+                    env=no_checkout_env,
+                )
+                # Apply patch.
+                _git_apply(patch_filename)
+                # Restore new saved index.
+                cmd_output_b('git', 'read-tree', new_tree, env=no_checkout_env)
+
             logger.info(f'Restored changes from {patch_filename}.')
     else:
         # There weren't any staged files so we don't need to do anything