Patchwork [2,of,2] ui: proxy protect/restorestdio() calls to update internal flag

login
register
mail settings
Submitter Yuya Nishihara
Date Jan. 12, 2019, 10:50 a.m.
Message ID <8e956cead2023a3afdfb.1547290204@mimosa>
Download mbox | patch
Permalink /patch/37691/
State Accepted
Headers show

Comments

Yuya Nishihara - Jan. 12, 2019, 10:50 a.m.
# HG changeset patch
# User Yuya Nishihara <yuya@tcha.org>
# Date 1537965712 -32400
#      Wed Sep 26 21:41:52 2018 +0900
# Node ID 8e956cead2023a3afdfb3f06a08b24290dffc244
# Parent  c5a71f1cdd3278159500360435863ee01c842cf4
ui: proxy protect/restorestdio() calls to update internal flag

It should be better to manage the redirection flag solely by the ui class.
Augie Fackler - Jan. 18, 2019, 9:48 a.m.
On Sat, Jan 12, 2019 at 07:50:04PM +0900, Yuya Nishihara wrote:
> # HG changeset patch
> # User Yuya Nishihara <yuya@tcha.org>
> # Date 1537965712 -32400
> #      Wed Sep 26 21:41:52 2018 +0900
> # Node ID 8e956cead2023a3afdfb3f06a08b24290dffc244
> # Parent  c5a71f1cdd3278159500360435863ee01c842cf4
> ui: proxy protect/restorestdio() calls to update internal flag

queued, thanks

Patch

diff --git a/mercurial/chgserver.py b/mercurial/chgserver.py
--- a/mercurial/chgserver.py
+++ b/mercurial/chgserver.py
@@ -200,7 +200,7 @@  def _newchgui(srcui, csystem, attachio):
         def _runsystem(self, cmd, environ, cwd, out):
             # fallback to the original system method if
             #  a. the output stream is not stdout (e.g. stderr, cStringIO),
-            #  b. or stdout is redirected by protectstdio(),
+            #  b. or stdout is redirected by protectfinout(),
             # because the chg client is not aware of these situations and
             # will behave differently (i.e. write to stdout).
             if (out is not self.fout
diff --git a/mercurial/ui.py b/mercurial/ui.py
--- a/mercurial/ui.py
+++ b/mercurial/ui.py
@@ -1080,14 +1080,38 @@  class ui(object):
             return False
         return procutil.isatty(fh)
 
+    def protectfinout(self):
+        """Duplicate ui streams and redirect original if they are stdio
+
+        Returns (fin, fout) which point to the original ui fds, but may be
+        copy of them. The returned streams can be considered "owned" in that
+        print(), exec(), etc. never reach to them.
+        """
+        if self._finoutredirected:
+            # if already redirected, protectstdio() would just create another
+            # nullfd pair, which is equivalent to returning self._fin/_fout.
+            return self._fin, self._fout
+        fin, fout = procutil.protectstdio(self._fin, self._fout)
+        self._finoutredirected = (fin, fout) != (self._fin, self._fout)
+        return fin, fout
+
+    def restorefinout(self, fin, fout):
+        """Restore ui streams from possibly duplicated (fin, fout)"""
+        if (fin, fout) == (self._fin, self._fout):
+            return
+        procutil.restorestdio(self._fin, self._fout, fin, fout)
+        # protectfinout() won't create more than one duplicated streams,
+        # so we can just turn the redirection flag off.
+        self._finoutredirected = False
+
     @contextlib.contextmanager
     def protectedfinout(self):
         """Run code block with protected standard streams"""
-        fin, fout = procutil.protectstdio(self._fin, self._fout)
+        fin, fout = self.protectfinout()
         try:
             yield fin, fout
         finally:
-            procutil.restorestdio(self._fin, self._fout, fin, fout)
+            self.restorefinout(fin, fout)
 
     def disablepager(self):
         self._disablepager = True
diff --git a/mercurial/wireprotoserver.py b/mercurial/wireprotoserver.py
--- a/mercurial/wireprotoserver.py
+++ b/mercurial/wireprotoserver.py
@@ -24,7 +24,6 @@  from . import (
 from .utils import (
     cborutil,
     interfaceutil,
-    procutil,
 )
 
 stringio = util.stringio
@@ -782,9 +781,7 @@  class sshserver(object):
     def __init__(self, ui, repo, logfh=None):
         self._ui = ui
         self._repo = repo
-        self._fin, self._fout = procutil.protectstdio(ui.fin, ui.fout)
-        # TODO: manage the redirection flag internally by ui
-        ui._finoutredirected = (self._fin, self._fout) != (ui.fin, ui.fout)
+        self._fin, self._fout = ui.protectfinout()
 
         # Log write I/O to stdout and stderr if configured.
         if logfh:
@@ -795,8 +792,7 @@  class sshserver(object):
 
     def serve_forever(self):
         self.serveuntil(threading.Event())
-        procutil.restorestdio(self._ui.fin, self._ui.fout,
-                              self._fin, self._fout)
+        self._ui.restorefinout(self._fin, self._fout)
         sys.exit(0)
 
     def serveuntil(self, ev):
diff --git a/tests/test-sshserver.py b/tests/test-sshserver.py
--- a/tests/test-sshserver.py
+++ b/tests/test-sshserver.py
@@ -47,6 +47,12 @@  class mockui(object):
         self.fout = io.BytesIO()
         self.ferr = io.BytesIO()
 
+    def protectfinout(self):
+        return self.fin, self.fout
+
+    def restorefinout(self, fin, fout):
+        pass
+
 if __name__ == '__main__':
     # Don't call into msvcrt to set BytesIO to binary mode
     procutil.setbinary = lambda fp: True