Don't add end-of-file newline while trimming whitespace

pre_commit_hooks/trailing_whitespace_fixer.py

@@ -19,14 +19,19 @@ def _fix_file(filename, is_markdown):
 
 
 def _process_line(line, is_markdown):
+    if line[-2:] == b'\r\n':
+        eol = b'\r\n'
+    elif line[-1:] == b'\n':
+        eol = b'\n'
+    else:
+        eol = b''
     # 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):
+def main(argv=None):
     parser = argparse.ArgumentParser()
     parser.add_argument(
         '--no-markdown-linebreak-ext',
@@ -77,4 +82,4 @@ def fix_trailing_whitespace(argv=None):
 
 
 if __name__ == '__main__':
-    sys.exit(fix_trailing_whitespace())
+    sys.exit(main())

tests/trailing_whitespace_fixer_test.py

@@ -3,7 +3,7 @@
 
 import pytest
 
-from pre_commit_hooks.trailing_whitespace_fixer import fix_trailing_whitespace
+from pre_commit_hooks.trailing_whitespace_fixer import main
 
 
 @pytest.mark.parametrize(
@@ -16,29 +16,37 @@
 def test_fixes_trailing_whitespace(input_s, expected, tmpdir):
     path = tmpdir.join('file.txt')
     path.write(input_s)
-    assert fix_trailing_whitespace((path.strpath,)) == 1
+    assert main((path.strpath,)) == 1
     assert path.read() == expected
 
 
+def test_ok_no_newline_end_of_file(tmpdir):
+    filename = tmpdir.join('f')
+    filename.write_binary(b'foo\nbar')
+    ret = main((filename.strpath,))
+    assert filename.read_binary() == b'foo\nbar'
+    assert ret == 0
+
+
 def test_ok_with_dos_line_endings(tmpdir):
     filename = tmpdir.join('f')
     filename.write_binary(b'foo\r\nbar\r\nbaz\r\n')
-    ret = fix_trailing_whitespace((filename.strpath,))
+    ret = main((filename.strpath,))
     assert filename.read_binary() == b'foo\r\nbar\r\nbaz\r\n'
     assert ret == 0
 
 
 def test_markdown_ok(tmpdir):
     filename = tmpdir.join('foo.md')
     filename.write_binary(b'foo  \n')
-    ret = fix_trailing_whitespace((filename.strpath,))
+    ret = main((filename.strpath,))
     assert filename.read_binary() == b'foo  \n'
     assert ret == 0
 
 
 # filename, expected input, expected output
 MD_TESTS_1 = (
-    ('foo.md', 'foo  \nbar \n  ', 'foo  \nbar\n\n'),
+    ('foo.md', 'foo  \nbar \n  ', 'foo  \nbar\n'),
     ('bar.Markdown', 'bar   \nbaz\t\n\t\n', 'bar  \nbaz\n\n'),
     ('.md', 'baz   \nquux  \t\n\t\n', 'baz\nquux\n\n'),
     ('txt', 'foo   \nbaz \n\t\n', 'foo\nbaz\n\n'),
@@ -49,7 +57,7 @@ def test_markdown_ok(tmpdir):
 def test_fixes_trailing_markdown_whitespace(filename, input_s, output, tmpdir):
     path = tmpdir.join(filename)
     path.write(input_s)
-    ret = fix_trailing_whitespace([path.strpath])
+    ret = main([path.strpath])
     assert ret == 1
     assert path.read() == output
 
@@ -68,16 +76,14 @@ def test_fixes_trailing_markdown_whitespace(filename, input_s, output, tmpdir):
 def test_markdown_linebreak_ext_opt(filename, input_s, output, tmpdir):
     path = tmpdir.join(filename)
     path.write(input_s)
-    ret = fix_trailing_whitespace((
-        '--markdown-linebreak-ext=TxT', path.strpath,
-    ))
+    ret = main(('--markdown-linebreak-ext=TxT', path.strpath))
     assert ret == 1
     assert path.read() == output
 
 
 # filename, expected input, expected output
 MD_TESTS_3 = (
-    ('foo.baz', 'foo  \nbar \n  ', 'foo  \nbar\n\n'),
+    ('foo.baz', 'foo  \nbar \n  ', 'foo  \nbar\n'),
     ('bar', 'bar   \nbaz\t\n\t\n', 'bar  \nbaz\n\n'),
 )
 
@@ -87,17 +93,15 @@ def test_markdown_linebreak_ext_opt_all(filename, input_s, output, tmpdir):
     path = tmpdir.join(filename)
     path.write(input_s)
     # need to make sure filename is not treated as argument to option
-    ret = fix_trailing_whitespace([
-        '--markdown-linebreak-ext=*', path.strpath,
-    ])
+    ret = main(('--markdown-linebreak-ext=*', path.strpath))
     assert ret == 1
     assert path.read() == output
 
 
 @pytest.mark.parametrize(('arg'), ('--', 'a.b', 'a/b'))
 def test_markdown_linebreak_ext_badopt(arg):
     with pytest.raises(SystemExit) as excinfo:
-        fix_trailing_whitespace(['--markdown-linebreak-ext', arg])
+        main(['--markdown-linebreak-ext', arg])
     assert excinfo.value.code == 2
 
 
@@ -112,19 +116,15 @@ def test_markdown_linebreak_ext_badopt(arg):
 def test_no_markdown_linebreak_ext_opt(filename, input_s, output, tmpdir):
     path = tmpdir.join(filename)
     path.write(input_s)
-    ret = fix_trailing_whitespace(['--no-markdown-linebreak-ext', path.strpath])
+    ret = main(['--no-markdown-linebreak-ext', path.strpath])
     assert ret == 1
     assert path.read() == output
 
 
-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])
+    ret = main([path.strpath])
     assert ret == 1
     assert path.size() == (len(non_utf8_bytes_content) - 1)