|
10 | 10 | (trim trailing whitespace, ensure final newline). |
11 | 11 | """ |
12 | 12 |
|
13 | | -from fileinput import input as file_input |
14 | | -from sys import stdin, stdout, modules |
| 13 | +from pathlib import Path |
| 14 | +from sys import argv, stdin, stdout, modules |
15 | 15 | from stdisplay.stdisplay import stdisplay |
16 | 16 |
|
17 | 17 |
|
18 | 18 | def main() -> None: |
19 | 19 | """ |
20 | | - main |
| 20 | + Safely print stdin or file to stdout with tweaks |
| 21 | + (trim trailing whitespace, ensure final newline). |
21 | 22 | """ |
22 | 23 | # https://github.com/pytest-dev/pytest/issues/4843 |
23 | | - if "pytest" not in modules: |
| 24 | + if "pytest" not in modules and stdin is not None: |
24 | 25 | stdin.reconfigure(errors="ignore") # type: ignore |
25 | | - |
26 | | - for line in file_input(encoding="ascii", errors="replace"): |
27 | | - stdout.write(stdisplay(line).rstrip() + "\n") |
28 | | - |
| 26 | + if len(argv) == 1: |
| 27 | + if stdin is not None: |
| 28 | + for untrusted_line in stdin: |
| 29 | + stdout.write(stdisplay(untrusted_line).rstrip() + "\n") |
| 30 | + stdout.flush() |
| 31 | + return |
| 32 | + for untrusted_arg in argv[1:]: |
| 33 | + if untrusted_arg == "-": |
| 34 | + if stdin is not None: |
| 35 | + for untrusted_line in stdin: |
| 36 | + stdout.write(stdisplay(untrusted_line).rstrip() + "\n") |
| 37 | + else: |
| 38 | + path = Path(untrusted_arg) |
| 39 | + untrusted_text = path.read_text(encoding="ascii", errors="replace") |
| 40 | + stdout.write(stdisplay(untrusted_text).rstrip() + "\n") |
29 | 41 | stdout.flush() |
30 | 42 |
|
31 | 43 |
|
|
0 commit comments