From b683657f2363ab048d4f7aaace9a73c9af3a9a69 Mon Sep 17 00:00:00 2001
From: Anthony Sottile <asottile@umich.edu>
Date: Thu, 19 Mar 2020 18:52:03 -0700
Subject: [PATCH] Support `babi -` for reading from stdin

Resolves #42
---
 babi/file.py                 | 10 ++++++++--
 babi/main.py                 | 26 +++++++++++++++++++++-----
 tests/features/stdin_test.py | 22 ++++++++++++++++++++++
 3 files changed, 51 insertions(+), 7 deletions(-)
 create mode 100644 tests/features/stdin_test.py

diff --git a/babi/file.py b/babi/file.py
index 9800f48..4dfa586 100644
--- a/babi/file.py
+++ b/babi/file.py
@@ -226,11 +226,17 @@ class File:
         self.selection = Selection()
         self._file_hls: Tuple[FileHL, ...] = ()
 
-    def ensure_loaded(self, status: Status) -> None:
+    def ensure_loaded(self, status: Status, stdin: str) -> None:
         if self.lines:
             return
 
-        if self.filename is not None and os.path.isfile(self.filename):
+        if self.filename == '-':
+            status.update('(from stdin)')
+            self.filename = None
+            self.modified = True
+            sio = io.StringIO(stdin)
+            self.lines, self.nl, mixed, self.sha256 = get_lines(sio)
+        elif self.filename is not None and os.path.isfile(self.filename):
             with open(self.filename, newline='') as f:
                 self.lines, self.nl, mixed, self.sha256 = get_lines(f)
         else:
diff --git a/babi/main.py b/babi/main.py
index 11a95d7..236a691 100644
--- a/babi/main.py
+++ b/babi/main.py
@@ -1,5 +1,7 @@
 import argparse
 import curses
+import os
+import sys
 from typing import Optional
 from typing import Sequence
 
@@ -9,9 +11,11 @@ from babi.screen import EditResult
 from babi.screen import make_stdscr
 from babi.screen import Screen
 
+CONSOLE = 'CONIN$' if sys.platform == 'win32' else '/dev/tty'
 
-def _edit(screen: Screen) -> EditResult:
-    screen.file.ensure_loaded(screen.status)
+
+def _edit(screen: Screen, stdin: str) -> EditResult:
+    screen.file.ensure_loaded(screen.status, stdin)
 
     while True:
         screen.status.tick(screen.margin)
@@ -32,13 +36,17 @@ def _edit(screen: Screen) -> EditResult:
             screen.status.update(f'unknown key: {key}')
 
 
-def c_main(stdscr: 'curses._CursesWindow', args: argparse.Namespace) -> int:
+def c_main(
+        stdscr: 'curses._CursesWindow',
+        args: argparse.Namespace,
+        stdin: str,
+) -> int:
     with perf_log(args.perf_log) as perf:
         screen = Screen(stdscr, args.filenames or [None], perf)
         with screen.history.save():
             while screen.files:
                 screen.i = screen.i % len(screen.files)
-                res = _edit(screen)
+                res = _edit(screen, stdin)
                 if res == EditResult.EXIT:
                     del screen.files[screen.i]
                     screen.status.clear()
@@ -59,8 +67,16 @@ def main(argv: Optional[Sequence[str]] = None) -> int:
     parser.add_argument('--perf-log')
     args = parser.parse_args(argv)
 
+    if '-' in args.filenames:
+        print('reading stdin...', file=sys.stderr)
+        stdin = sys.stdin.read()
+        tty = os.open(CONSOLE, os.O_RDONLY)
+        os.dup2(tty, sys.stdin.fileno())
+    else:
+        stdin = ''
+
     with make_stdscr() as stdscr:
-        return c_main(stdscr, args)
+        return c_main(stdscr, args, stdin)
 
 
 if __name__ == '__main__':
diff --git a/tests/features/stdin_test.py b/tests/features/stdin_test.py
new file mode 100644
index 0000000..169388c
--- /dev/null
+++ b/tests/features/stdin_test.py
@@ -0,0 +1,22 @@
+import shlex
+import sys
+
+from babi.screen import VERSION_STR
+from testing.runner import PrintsErrorRunner
+
+
+def test_open_from_stdin():
+    with PrintsErrorRunner('env', 'TERM=screen', 'bash', '--norc') as h:
+        cmd = (sys.executable, '-mcoverage', 'run', '-m', 'babi', '-')
+        babi_cmd = ' '.join(shlex.quote(part) for part in cmd)
+        h.press_and_enter(fr"echo $'hello\nworld' | {babi_cmd}")
+
+        h.await_text(VERSION_STR, timeout=2)
+        h.await_text('<<new file>> *')
+        h.await_text('hello\nworld')
+
+        h.press('^X')
+        h.press('n')
+        h.await_text_missing('<<new file>>')
+        h.press_and_enter('exit')
+        h.await_exit()