improve sub process handling: return child exit code

Wrap the management of the sub process in a new class ForkPty. This
reduces the length and number of variables in the main function. Turn it
into a context manager to allow closing the masterfd and actually
waiting for the forked child (rather than reparenting it to init on
exit). This also allows propagating the exit code. Document the new exit
semantics.

2 files changed, 101 insertions, 71 deletions

diff --git a/tcvt.1 b/tcvt.1
index 8d5af3c..cca5e11 100644
--- a/tcvt.1
+++ b/tcvt.1
@@ -37,6 +37,9 @@ Signal end of options.
 .El
 .Pp
 Any non-option argument signals the end of options for \fBtcvt\fP.
+.Sh EXIT STATUS
+The exit status of the wrapped program is propagated.
+When \fBtcvt\fP produces an error of its own, 255 is returned.
 .Sh BUGS
 Not all \fBANSI\fP terminal features are emulated.
 If you experience rendering issues, you can export the environment variable \fBTCVT_DEVEL\fP.
diff --git a/tcvt.py b/tcvt.py
index 97df94c..edb102f 100755
--- a/tcvt.py
+++ b/tcvt.py
@@ -649,6 +649,46 @@ def set_cloexec(fd):
     flags |= fcntl.FD_CLOEXEC
     fcntl.fcntl(fd, fcntl.F_SETFD, flags)
 
+class ExecutionError(Exception):
+    pass
+
+class ForkPty(object):
+    def __init__(self, argv, environ={}):
+        self.argv = argv
+        self.environ = environ
+        self.pid = -1
+        self.masterfd = -1
+        self.exitcode = 255
+
+    def __enter__(self):
+        assert self.pid == -1
+        assert self.masterfd == -1
+        errpiper, errpipew = os.pipe()
+        set_cloexec(errpipew)
+        self.pid, self.masterfd = pty.fork()
+        if self.pid == 0: # child
+            os.close(errpiper)
+            os.environ.update(self.environ)
+            try:
+                os.execvp(self.argv[0], self.argv)
+            except OSError as err:
+                os.write(errpipew, "exec failed: %s" % (err,))
+            sys.exit(255)
+
+        os.close(errpipew)
+        data = os.read(errpiper, 1024)
+        os.close(errpiper)
+        if data:
+            raise ExecutionError(data)
+
+        return self.masterfd
+
+    def __exit__(self, *_):
+        os.close(self.masterfd)
+        status = os.waitpid(self.pid, 0)[1]
+        if status & 0xff == 0: # not killed by a signal
+            self.exitcode = status >> 8
+
 def main():
     parser = optparse.OptionParser()
     parser.disable_interspersed_args()
@@ -660,79 +700,66 @@ def main():
     options, args = parser.parse_args()
     keymapping, acsc = compute_keymap(symbolic_keymapping)
 
-    errpiper, errpipew = os.pipe()
-    set_cloexec(errpipew)
-    pid, masterfd = pty.fork()
-    if pid == 0: # child
-        os.close(errpiper)
-        os.environ["TERM"] = "ansi"
-        try:
-            if len(args) < 1:
-                os.execvp(os.environ["SHELL"], [os.environ["SHELL"]])
-            else:
-                os.execvp(args[0], args)
-        except OSError as err:
-            os.write(errpipew, "exec failed: %s" % (err,))
-        sys.exit(1)
-
-    os.close(errpipew)
-    data = os.read(errpiper, 1024)
-    os.close(errpiper)
-    if data:
-        print(data)
-        sys.exit(1)
-
-    with Terminal(acsc, options.columns, reverse=options.reverse) as t:
-        t.resizepty(masterfd)
-        refreshpending = None
-        while True:
-            try:
-                res, _, _ = select.select([0, masterfd], [], [],
-                                          refreshpending and 0)
-            except select.error as err:
-                if err.args[0] == errno.EINTR:
-                    t.resized()
-                    t.resizepty(masterfd)
-                    continue
-                raise
-            if 0 in res:
+    process = ForkPty(args or [os.environ["SHELL"]], dict(TERM="ansi"))
+    try:
+        with process as masterfd:
+            with Terminal(acsc, options.columns, reverse=options.reverse) as t:
+                t.resizepty(masterfd)
+                refreshpending = None
                 while True:
-                    key = t.realscreen.getch()
-                    if key == -1:
-                        break
-                    if key == 0xb3:
-                        t.switchmode()
-                        t.resizepty(masterfd)
-                    elif key in keymapping:
-                        os.write(masterfd, keymapping[key])
-                    elif key <= 0xff:
-                        os.write(masterfd, struct.pack("B", key))
-                    else:
-                        if "TCVT_DEVEL" in os.environ:
-                            raise ValueError("getch returned %d" % key)
-            elif masterfd in res:
-                try:
-                    data = os.read(masterfd, 1024)
-                except OSError:
-                    break
-                if not data:
-                    break
-                for char in bytearray(data):
-                    if "TCVT_DEVEL" in os.environ:
-                        t.feed(char)
-                    else:
+                    try:
+                        res, _, _ = select.select([0, masterfd], [], [],
+                                                  refreshpending and 0)
+                    except select.error as err:
+                        if err.args[0] == errno.EINTR:
+                            t.resized()
+                            t.resizepty(masterfd)
+                            continue
+                        raise
+                    if 0 in res:
+                        while True:
+                            key = t.realscreen.getch()
+                            if key == -1:
+                                break
+                            if key == 0xb3:
+                                t.switchmode()
+                                t.resizepty(masterfd)
+                            elif key in keymapping:
+                                os.write(masterfd, keymapping[key])
+                            elif key <= 0xff:
+                                os.write(masterfd, struct.pack("B", key))
+                            else:
+                                if "TCVT_DEVEL" in os.environ:
+                                    raise ValueError("getch returned %d" % key)
+                    elif masterfd in res:
                         try:
-                            t.feed(char)
-                        except ValueError:
-                            t.feed_reset()
-                if refreshpending is None:
-                    refreshpending = time.time() + 0.1
-            elif refreshpending is not None:
-                t.screen.refresh()
-                refreshpending = None
-            if refreshpending is not None and refreshpending < time.time():
-                t.screen.refresh()
-                refreshpending = None
+                            data = os.read(masterfd, 1024)
+                        except OSError:
+                            break
+                        if not data:
+                            break
+                        for char in bytearray(data):
+                            if "TCVT_DEVEL" in os.environ:
+                                t.feed(char)
+                            else:
+                                try:
+                                    t.feed(char)
+                                except ValueError:
+                                    t.feed_reset()
+                        if refreshpending is None:
+                            refreshpending = time.time() + 0.1
+                    elif refreshpending is not None:
+                        t.screen.refresh()
+                        refreshpending = None
+                    if refreshpending is not None and \
+                            refreshpending < time.time():
+                        t.screen.refresh()
+                        refreshpending = None
+    except ExecutionError as err:
+        print(str(err))
+        sys.exit(255)
+    else:
+        sys.exit(process.exitcode)
 
 if __name__ == '__main__':
     main()