1 files changed, 90 insertions, 0 deletions
diff --git a/pre_commit/staged_files_only.py b/pre_commit/staged_files_only.py
new file mode 100644
index 0000000..09d323d
--- /dev/null
+++ b/pre_commit/staged_files_only.py
@@ -0,0 +1,90 @@
+import contextlib
+import logging
+import os.path
+import time
+from typing import Generator
+
+from pre_commit import git
+from pre_commit.util import CalledProcessError
+from pre_commit.util import cmd_output
+from pre_commit.util import cmd_output_b
+from pre_commit.xargs import xargs
+
+
+logger = logging.getLogger('pre_commit')
+
+
+def _git_apply(patch: str) -> None:
+    args = ('apply', '--whitespace=nowarn', patch)
+    try:
+        cmd_output_b('git', *args)
+    except CalledProcessError:
+        # Retry with autocrlf=false -- see #570
+        cmd_output_b('git', '-c', 'core.autocrlf=false', *args)
+
+
+@contextlib.contextmanager
+def _intent_to_add_cleared() -> Generator[None, None, None]:
+    intent_to_add = git.intent_to_add_files()
+    if intent_to_add:
+        logger.warning('Unstaged intent-to-add files detected.')
+
+        xargs(('git', 'rm', '--cached', '--'), intent_to_add)
+        try:
+            yield
+        finally:
+            xargs(('git', 'add', '--intent-to-add', '--'), intent_to_add)
+    else:
+        yield
+
+
+@contextlib.contextmanager
+def _unstaged_changes_cleared(patch_dir: str) -> Generator[None, None, None]:
+    tree = cmd_output('git', 'write-tree')[1].strip()
+    retcode, diff_stdout_binary, _ = cmd_output_b(
+        'git', 'diff-index', '--ignore-submodules', '--binary',
+        '--exit-code', '--no-color', '--no-ext-diff', tree, '--',
+        retcode=None,
+    )
+    if retcode and diff_stdout_binary.strip():
+        patch_filename = f'patch{int(time.time())}'
+        patch_filename = os.path.join(patch_dir, patch_filename)
+        logger.warning('Unstaged files detected.')
+        logger.info(f'Stashing unstaged files to {patch_filename}.')
+        # Save the current unstaged changes as a patch
+        os.makedirs(patch_dir, exist_ok=True)
+        with open(patch_filename, 'wb') as patch_file:
+            patch_file.write(diff_stdout_binary)
+
+        # Clear the working directory of unstaged changes
+        cmd_output_b('git', 'checkout', '--', '.')
+        try:
+            yield
+        finally:
+            # Try to apply the patch we saved
+            try:
+                _git_apply(patch_filename)
+            except CalledProcessError:
+                logger.warning(
+                    'Stashed changes conflicted with hook auto-fixes... '
+                    'Rolling back fixes...',
+                )
+                # We failed to apply the patch, presumably due to fixes made
+                # by hooks.
+                # Roll back the changes made by hooks.
+                cmd_output_b('git', 'checkout', '--', '.')
+                _git_apply(patch_filename)
+            logger.info(f'Restored changes from {patch_filename}.')
+    else:
+        # There weren't any staged files so we don't need to do anything
+        # special
+        yield
+
+
+@contextlib.contextmanager
+def staged_files_only(patch_dir: str) -> Generator[None, None, None]:
+    """Clear any unstaged changes from the git working directory inside this
+    context.
+    """
+    with _intent_to_add_cleared(), _unstaged_changes_cleared(patch_dir):
+        yield