Merge pull request #135 from Lucas-C/master

trailing-whitespace hook: restoring original file in case of failure - fixes #134

pre_commit_hooks/trailing_whitespace_fixer.py

@@ -1,24 +1,27 @@
 from __future__ import print_function
 
 import argparse
-import fileinput
 import os
 import sys
 
 from pre_commit_hooks.util import cmd_output
 
 
-def _fix_file(filename, markdown=False):
+def _fix_file(filename, is_markdown):
-    for line in fileinput.input([filename], inplace=True):
+    with open(filename, mode='rb') as file_processed:
-        # preserve trailing two-space for non-blank lines in markdown files
+        lines = file_processed.readlines()
-        if markdown and (not line.isspace()) and (line.endswith("  \n")):
+    lines = [_process_line(line, is_markdown) for line in lines]
-            line = line.rstrip(' \n')
+    with open(filename, mode='wb') as file_processed:
-            # only preserve if there are no trailing tabs or unusual whitespace
+        for line in lines:
-            if not line[-1].isspace():
+            file_processed.write(line)
-                print(line + "  ")
-                continue
 
-        print(line.rstrip())
+
+def _process_line(line, is_markdown):
+    # preserve trailing two-space for non-blank lines in markdown files
+    eol = b'\r\n' if line[-2:] == b'\r\n' else b'\n'
+    if is_markdown and (not line.isspace()) and line.endswith(b'  ' + eol):
+        return line.rstrip() + b'  ' + eol
+    return line.rstrip() + eol
 
 
 def fix_trailing_whitespace(argv=None):
@@ -64,14 +67,13 @@ def fix_trailing_whitespace(argv=None):
                 .format(ext)
             )
 
-    if bad_whitespace_files:
+    return_code = 0
     for bad_whitespace_file in bad_whitespace_files:
         print('Fixing {0}'.format(bad_whitespace_file))
         _, extension = os.path.splitext(bad_whitespace_file.lower())
         _fix_file(bad_whitespace_file, all_markdown or extension in md_exts)
-        return 1
+        return_code = 1
-    else:
+    return return_code
-        return 0
 
 
 if __name__ == '__main__':

tests/trailing_whitespace_fixer_test.py

@@ -103,3 +103,12 @@ def test_no_markdown_linebreak_ext_opt(filename, input_s, output, tmpdir):
 
 def test_returns_zero_for_no_changes():
     assert fix_trailing_whitespace([__file__]) == 0
+
+
+def test_preserve_non_utf8_file(tmpdir):
+    non_utf8_bytes_content = b'<a>\xe9 \n</a>\n'
+    path = tmpdir.join('file.txt')
+    path.write_binary(non_utf8_bytes_content)
+    ret = fix_trailing_whitespace([path.strpath])
+    assert ret == 1
+    assert path.size() == (len(non_utf8_bytes_content) - 1)