diff --git a/README.md b/README.md index 584b53d..e579976 100644 --- a/README.md +++ b/README.md @@ -111,6 +111,8 @@ Add this to your `.pre-commit-config.yaml` use `args: [--markdown-linebreak-ext=md]` (or other extensions used by your markdownfiles). If for some reason you want to treat all files as markdown, use `--markdown-linebreak-ext=*`. + - By default, this hook trims all whitespace from the ends of lines. + To specify a custom set of characters to trim instead, use `args: [--chars,""]`. ### Deprecated / replaced hooks diff --git a/pre_commit_hooks/trailing_whitespace_fixer.py b/pre_commit_hooks/trailing_whitespace_fixer.py index 2ccc003..a21b54f 100644 --- a/pre_commit_hooks/trailing_whitespace_fixer.py +++ b/pre_commit_hooks/trailing_whitespace_fixer.py @@ -7,10 +7,11 @@ from typing import Optional from typing import Sequence -def _fix_file(filename, is_markdown): # type: (str, bool) -> bool +def _fix_file(filename, is_markdown, chars): + # type: (str, bool, Optional[bytes]) -> bool with open(filename, mode='rb') as file_processed: lines = file_processed.readlines() - newlines = [_process_line(line, is_markdown) for line in lines] + newlines = [_process_line(line, is_markdown, chars) for line in lines] if newlines != lines: with open(filename, mode='wb') as file_processed: for line in newlines: @@ -20,17 +21,20 @@ def _fix_file(filename, is_markdown): # type: (str, bool) -> bool return False -def _process_line(line, is_markdown): # type: (bytes, bool) -> bytes +def _process_line(line, is_markdown, chars): + # type: (bytes, bool, Optional[bytes]) -> bytes if line[-2:] == b'\r\n': eol = b'\r\n' + line = line[:-2] elif line[-1:] == b'\n': eol = b'\n' + line = line[:-1] else: eol = b'' # preserve trailing two-space for non-blank lines in markdown files - if is_markdown and (not line.isspace()) and line.endswith(b' ' + eol): - return line.rstrip() + b' ' + eol - return line.rstrip() + eol + if is_markdown and (not line.isspace()) and line.endswith(b' '): + return line[:-2].rstrip(chars) + b' ' + eol + return line.rstrip(chars) + eol def main(argv=None): # type: (Optional[Sequence[str]]) -> int @@ -50,6 +54,13 @@ def main(argv=None): # type: (Optional[Sequence[str]]) -> int 'default: %(default)s' ), ) + parser.add_argument( + '--chars', + help=( + 'The set of characters to strip from the end of lines. ' + 'Defaults to all whitespace characters.' + ), + ) parser.add_argument('filenames', nargs='*', help='Filenames to fix') args = parser.parse_args(argv) @@ -73,12 +84,12 @@ def main(argv=None): # type: (Optional[Sequence[str]]) -> int " (probably filename; use '--markdown-linebreak-ext=EXT')" .format(ext), ) - + chars = None if args.chars is None else args.chars.encode('utf-8') return_code = 0 for filename in args.filenames: _, extension = os.path.splitext(filename.lower()) md = all_markdown or extension in md_exts - if _fix_file(filename, md): + if _fix_file(filename, md, chars): print('Fixing {}'.format(filename)) return_code = 1 return return_code diff --git a/tests/trailing_whitespace_fixer_test.py b/tests/trailing_whitespace_fixer_test.py index 82c9b6d..97f9aef 100644 --- a/tests/trailing_whitespace_fixer_test.py +++ b/tests/trailing_whitespace_fixer_test.py @@ -78,3 +78,27 @@ def test_preserve_non_utf8_file(tmpdir): ret = main([path.strpath]) assert ret == 1 assert path.size() == (len(non_utf8_bytes_content) - 1) + + +def test_custom_charset_change(tmpdir): + # strip spaces only, no tabs + path = tmpdir.join('file.txt') + path.write('\ta \t \n') + ret = main([path.strpath, '--chars', ' ']) + assert ret == 1 + assert path.read() == '\ta \t\n' + + +def test_custom_charset_no_change(tmpdir): + path = tmpdir.join('file.txt') + path.write('\ta \t\n') + ret = main([path.strpath, '--chars', ' ']) + assert ret == 0 + + +def test_markdown_with_custom_charset(tmpdir): + path = tmpdir.join('file.md') + path.write('\ta \t \n') + ret = main([path.strpath, '--chars', ' ', '--markdown-linebreak-ext', '*']) + assert ret == 1 + assert path.read() == '\ta \t \n'