D2081: wireprotoserver: add context manager mechanism for redirecting stdio

[indygreg (Gregory Szorc)]

indygreg updated this revision to Diff 5333.
indygreg edited the summary of this revision.

REPOSITORY
  rHG Mercurial

CHANGES SINCE LAST UPDATE
  https://phab.mercurial-scm.org/D2081?vs=5315&id=5333

REVISION DETAIL
  https://phab.mercurial-scm.org/D2081

AFFECTED FILES
  mercurial/wireproto.py
  mercurial/wireprotoserver.py

CHANGE DETAILS

diff --git a/mercurial/wireprotoserver.py b/mercurial/wireprotoserver.py
--- a/mercurial/wireprotoserver.py
+++ b/mercurial/wireprotoserver.py
@@ -8,6 +8,7 @@
 
 import abc
 import cgi
+import contextlib
 import struct
 import sys
 
@@ -74,6 +75,20 @@
         """
 
     @abc.abstractmethod
+    def mayberedirectstdio(self):
+        """Context manager to possibly redirect stdio.
+
+        The context manager yields a file-object like object that receives
+        stdout and stderr output when the context manager is active. Or it
+        yields ``None`` if no I/O redirection occurs.
+
+        The intent of this context manager is to capture stdio output
+        so it may be sent in the response. Some transports support streaming
+        stdio to the client in real time. For these transports, stdio output
+        won't be captured.
+        """
+
+    @abc.abstractmethod
     def redirect(self):
         """may setup interception for stdout and stderr
 
@@ -151,6 +166,21 @@
         for s in util.filechunkiter(self._req, limit=length):
             fp.write(s)
 
+    @contextlib.contextmanager
+    def mayberedirectstdio(self):
+        oldout = self._ui.fout
+        olderr = self._ui.ferr
+
+        out = util.stringio()
+
+        try:
+            self._ui.fout = out
+            self._ui.ferr = out
+            yield out
+        finally:
+            self._ui.fout = oldout
+            self._ui.ferr = olderr
+
     def redirect(self):
         self._oldio = self._ui.fout, self._ui.ferr
         self._ui.ferr = self._ui.fout = stringio()
@@ -393,6 +423,10 @@
             fpout.write(self._fin.read(count))
             count = int(self._fin.readline())
 
+    @contextlib.contextmanager
+    def mayberedirectstdio(self):
+        yield None
+
     def redirect(self):
         pass
 
diff --git a/mercurial/wireproto.py b/mercurial/wireproto.py
--- a/mercurial/wireproto.py
+++ b/mercurial/wireproto.py
@@ -978,20 +978,12 @@
     else:
         new = encoding.tolocal(new) # normal path
 
-    if util.safehasattr(proto, 'restore'):
-
-        proto.redirect()
-
+    with proto.mayberedirectstdio() as output:
         r = repo.pushkey(encoding.tolocal(namespace), encoding.tolocal(key),
                          encoding.tolocal(old), new) or False
 
-        output = proto.restore()
-
-        return '%s\n%s' % (int(r), output)
-
-    r = repo.pushkey(encoding.tolocal(namespace), encoding.tolocal(key),
-                     encoding.tolocal(old), new)
-    return '%s\n' % int(r)
+    output = output.getvalue() if output else ''
+    return '%s\n%s' % (int(r), output)
 
 @wireprotocommand('stream_out')
 def stream(repo, proto):



To: indygreg, #hg-reviewers
Cc: mercurial-devel