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