Anthony Sottile | 8f61529 | 2022-01-15 19:24:05 -0500 | [diff] [blame] | 1 | from __future__ import annotations |
| 2 | |
Anthony Sottile | 57f1533 | 2014-03-22 21:40:57 -0700 | [diff] [blame] | 3 | import argparse |
| 4 | import os |
Anthony Sottile | 030bfac | 2019-01-31 19:19:10 -0800 | [diff] [blame] | 5 | from typing import IO |
Anthony Sottile | 030bfac | 2019-01-31 19:19:10 -0800 | [diff] [blame] | 6 | from typing import Sequence |
Anthony Sottile | 57f1533 | 2014-03-22 21:40:57 -0700 | [diff] [blame] | 7 | |
Anthony Sottile | 57f1533 | 2014-03-22 21:40:57 -0700 | [diff] [blame] | 8 | |
Anthony Sottile | f5c42a0 | 2020-02-05 11:10:42 -0800 | [diff] [blame] | 9 | def fix_file(file_obj: IO[bytes]) -> int: |
Anthony Sottile | 57f1533 | 2014-03-22 21:40:57 -0700 | [diff] [blame] | 10 | # Test for newline at end of file |
| 11 | # Empty files will throw IOError here |
| 12 | try: |
| 13 | file_obj.seek(-1, os.SEEK_END) |
Anthony Sottile | f5c42a0 | 2020-02-05 11:10:42 -0800 | [diff] [blame] | 14 | except OSError: |
Anthony Sottile | 57f1533 | 2014-03-22 21:40:57 -0700 | [diff] [blame] | 15 | return 0 |
| 16 | last_character = file_obj.read(1) |
| 17 | # last_character will be '' for an empty file |
mtkennerly | e694a6c | 2018-10-13 19:44:02 -0400 | [diff] [blame] | 18 | if last_character not in {b'\n', b'\r'} and last_character != b'': |
Anthony Sottile | dd9a07d | 2015-01-19 16:43:10 -0800 | [diff] [blame] | 19 | # Needs this seek for windows, otherwise IOError |
| 20 | file_obj.seek(0, os.SEEK_END) |
Anthony Sottile | b80ca9e | 2014-04-13 22:09:14 -0700 | [diff] [blame] | 21 | file_obj.write(b'\n') |
Anthony Sottile | 57f1533 | 2014-03-22 21:40:57 -0700 | [diff] [blame] | 22 | return 1 |
| 23 | |
mtkennerly | e694a6c | 2018-10-13 19:44:02 -0400 | [diff] [blame] | 24 | while last_character in {b'\n', b'\r'}: |
Anthony Sottile | 57f1533 | 2014-03-22 21:40:57 -0700 | [diff] [blame] | 25 | # Deal with the beginning of the file |
| 26 | if file_obj.tell() == 1: |
| 27 | # If we've reached the beginning of the file and it is all |
| 28 | # linebreaks then we can make this file empty |
| 29 | file_obj.seek(0) |
| 30 | file_obj.truncate() |
| 31 | return 1 |
| 32 | |
| 33 | # Go back two bytes and read a character |
| 34 | file_obj.seek(-2, os.SEEK_CUR) |
| 35 | last_character = file_obj.read(1) |
| 36 | |
| 37 | # Our current position is at the end of the file just before any amount of |
mtkennerly | 2ab5832 | 2018-10-13 19:00:22 -0400 | [diff] [blame] | 38 | # newlines. If we find extraneous newlines, then backtrack and trim them. |
| 39 | position = file_obj.tell() |
| 40 | remaining = file_obj.read() |
mtkennerly | e694a6c | 2018-10-13 19:44:02 -0400 | [diff] [blame] | 41 | for sequence in (b'\n', b'\r\n', b'\r'): |
| 42 | if remaining == sequence: |
| 43 | return 0 |
| 44 | elif remaining.startswith(sequence): |
mtkennerly | 2ab5832 | 2018-10-13 19:00:22 -0400 | [diff] [blame] | 45 | file_obj.seek(position + len(sequence)) |
| 46 | file_obj.truncate() |
| 47 | return 1 |
Anthony Sottile | 57f1533 | 2014-03-22 21:40:57 -0700 | [diff] [blame] | 48 | |
| 49 | return 0 |
| 50 | |
| 51 | |
Anthony Sottile | 8f61529 | 2022-01-15 19:24:05 -0500 | [diff] [blame] | 52 | def main(argv: Sequence[str] | None = None) -> int: |
Anthony Sottile | 57f1533 | 2014-03-22 21:40:57 -0700 | [diff] [blame] | 53 | parser = argparse.ArgumentParser() |
| 54 | parser.add_argument('filenames', nargs='*', help='Filenames to fix') |
| 55 | args = parser.parse_args(argv) |
| 56 | |
| 57 | retv = 0 |
| 58 | |
| 59 | for filename in args.filenames: |
| 60 | # Read as binary so we can read byte-by-byte |
| 61 | with open(filename, 'rb+') as file_obj: |
| 62 | ret_for_file = fix_file(file_obj) |
| 63 | if ret_for_file: |
Anthony Sottile | f5c42a0 | 2020-02-05 11:10:42 -0800 | [diff] [blame] | 64 | print(f'Fixing {filename}') |
Anthony Sottile | 57f1533 | 2014-03-22 21:40:57 -0700 | [diff] [blame] | 65 | retv |= ret_for_file |
| 66 | |
| 67 | return retv |
| 68 | |
| 69 | |
| 70 | if __name__ == '__main__': |
Anthony Sottile | 39ab2ed | 2021-10-23 13:23:50 -0400 | [diff] [blame] | 71 | raise SystemExit(main()) |