Add staged_files_only context manager.

pre_commit/prefixed_command_runner.py

@@ -6,6 +6,9 @@ import subprocess
 
 class CalledProcessError(RuntimeError):
     def __init__(self, returncode, cmd, expected_returncode, output=None):
+        super(CalledProcessError, self).__init__(
+            returncode, cmd, expected_returncode, output,
+        )
         self.returncode = returncode
         self.cmd = cmd
         self.expected_returncode = expected_returncode
@@ -15,13 +18,13 @@ class CalledProcessError(RuntimeError):
         return (
             'Command: {0!r}\n'
             'Return code: {1}\n'
-            'Expected return code {2}\n',
+            'Expected return code: {2}\n'
             'Output: {3!r}\n'.format(
                 self.cmd,
                 self.returncode,
                 self.expected_returncode,
                 self.output,
-            ),
+            )
         )
 
 
@@ -48,15 +51,15 @@ class PrefixedCommandRunner(object):
             self.__makedirs(self.prefix_dir)
 
     def run(self, cmd, retcode=0, stdin=None, **kwargs):
+        popen_kwargs = {
+            'stdin': subprocess.PIPE,
+            'stdout': subprocess.PIPE,
+            'stderr': subprocess.PIPE,
+        }
+        popen_kwargs.update(kwargs)
         self._create_path_if_not_exists()
         replaced_cmd = _replace_cmd(cmd, prefix=self.prefix_dir)
-        proc = self.__popen(
-            replaced_cmd,
-            stdin=subprocess.PIPE,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            **kwargs
-        )
+        proc = self.__popen(replaced_cmd, **popen_kwargs)
         stdout, stderr = proc.communicate(stdin)
         returncode = proc.returncode
 

pre_commit/staged_files_only.py

@@ -0,0 +1,47 @@
+
+import contextlib
+import time
+
+from pre_commit.prefixed_command_runner import CalledProcessError
+
+
+@contextlib.contextmanager
+def staged_files_only(cmd_runner):
+    """Clear any unstaged changes from the git working directory inside this
+    context.
+
+    Args:
+        cmd_runner - PrefixedCommandRunner
+    """
+    # Determine if there are unstaged files
+    retcode, _, _ = cmd_runner.run(
+        ['git', 'diff-files', '--quiet'],
+        retcode=None,
+    )
+    if retcode:
+        # TODO: print a warning message that unstaged things are being stashed
+        # Save the current unstaged changes as a patch
+        # TODO: use a more unique patch filename
+        patch_filename = cmd_runner.path('patch{0}'.format(time.time()))
+        with open(patch_filename, 'w') as patch_file:
+            cmd_runner.run(['git', 'diff', '--binary'], stdout=patch_file)
+
+        # Clear the working directory of unstaged changes
+        cmd_runner.run(['git', 'checkout', '--', '.'])
+        try:
+            yield
+        finally:
+            # Try to apply the patch we saved
+            try:
+                cmd_runner.run(['git', 'apply', patch_filename])
+            except CalledProcessError:
+                # TOOD: print a warning about rolling back changes made by hooks
+                # We failed to apply the patch, presumably due to fixes made
+                # by hooks.
+                # Roll back the changes made by hooks.
+                cmd_runner.run(['git', 'checkout', '--', '.'])
+                cmd_runner.run(['git', 'apply', patch_filename])
+    else:
+        # There weren't any staged files so we don't need to do anything
+        # special
+        yield