changeset 37126:0216232f21ab

procutil: move protectio/restoreio from commandserver

Some variants of this will be useful for stdio-based servers such as
sshserver.

--- a/mercurial/commandserver.py	Sun Mar 25 11:30:59 2018 +0900
+++ b/mercurial/commandserver.py	Sun Mar 25 11:40:30 2018 +0900
@@ -306,35 +306,6 @@
 
         return 0
 
-def _protectio(uin, uout):
-    """Duplicate streams and redirect original to null if (uin, uout) are
-    stdio
-
-    Returns (fin, fout) which point to the original (uin, uout) fds, but
-    may be copy of (uin, uout). The returned streams can be considered
-    "owned" in that print(), exec(), etc. never reach to them.
-    """
-    uout.flush()
-    newfiles = []
-    nullfd = os.open(os.devnull, os.O_RDWR)
-    for f, sysf, mode in [(uin, procutil.stdin, r'rb'),
-                          (uout, procutil.stdout, r'wb')]:
-        if f is sysf:
-            newfd = os.dup(f.fileno())
-            os.dup2(nullfd, f.fileno())
-            f = os.fdopen(newfd, mode)
-        newfiles.append(f)
-    os.close(nullfd)
-    return tuple(newfiles)
-
-def _restoreio(uin, uout, fin, fout):
-    """Restore (uin, uout) streams from possibly duplicated (fin, fout)"""
-    uout.flush()
-    for f, uif in [(fin, uin), (fout, uout)]:
-        if f is not uif:
-            os.dup2(f.fileno(), uif.fileno())
-            f.close()
-
 class pipeservice(object):
     def __init__(self, ui, repo, opts):
         self.ui = ui
@@ -347,13 +318,13 @@
         ui = self.ui
         # redirect stdio to null device so that broken extensions or in-process
         # hooks will never cause corruption of channel protocol.
-        fin, fout = _protectio(ui.fin, ui.fout)
+        fin, fout = procutil.protectstdio(ui.fin, ui.fout)
         try:
             sv = server(ui, self.repo, fin, fout)
             return sv.serve()
         finally:
             sv.cleanup()
-            _restoreio(ui.fin, ui.fout, fin, fout)
+            procutil.restorestdio(ui.fin, ui.fout, fin, fout)
 
 def _initworkerprocess():
     # use a different process group from the master process, in order to:

--- a/mercurial/utils/procutil.py	Sun Mar 25 11:30:59 2018 +0900
+++ b/mercurial/utils/procutil.py	Sun Mar 25 11:40:30 2018 +0900
@@ -211,6 +211,35 @@
 def isstdout(f):
     return _testfileno(f, sys.__stdout__)
 
+def protectstdio(uin, uout):
+    """Duplicate streams and redirect original to null if (uin, uout) are
+    stdio
+
+    Returns (fin, fout) which point to the original (uin, uout) fds, but
+    may be copy of (uin, uout). The returned streams can be considered
+    "owned" in that print(), exec(), etc. never reach to them.
+    """
+    uout.flush()
+    newfiles = []
+    nullfd = os.open(os.devnull, os.O_RDWR)
+    for f, sysf, mode in [(uin, stdin, r'rb'),
+                          (uout, stdout, r'wb')]:
+        if f is sysf:
+            newfd = os.dup(f.fileno())
+            os.dup2(nullfd, f.fileno())
+            f = os.fdopen(newfd, mode)
+        newfiles.append(f)
+    os.close(nullfd)
+    return tuple(newfiles)
+
+def restorestdio(uin, uout, fin, fout):
+    """Restore (uin, uout) streams from possibly duplicated (fin, fout)"""
+    uout.flush()
+    for f, uif in [(fin, uin), (fout, uout)]:
+        if f is not uif:
+            os.dup2(f.fileno(), uif.fileno())
+            f.close()
+
 def shellenviron(environ=None):
     """return environ with optional override, useful for shelling out"""
     def py2shell(val):