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.

1 files changed, 98 insertions, 71 deletions

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()