Patchwork [3,of,5] procutil: introduce context-manager interface for protect/restorestdio

login
register
mail settings
Submitter Yuya Nishihara
Date March 25, 2018, 12:54 p.m.
Message ID <a315bc2695f7faff23e7.1521982454@mimosa>
Download mbox | patch
Permalink /patch/29851/
State Accepted
Headers show

Comments

Yuya Nishihara - March 25, 2018, 12:54 p.m.
# HG changeset patch
# User Yuya Nishihara <yuya@tcha.org>
# Date 1521946685 -32400
#      Sun Mar 25 11:58:05 2018 +0900
# Node ID a315bc2695f7faff23e71acea37d2c4a13961b14
# Parent  d3c212ea1f34ecff5ae12beceea0f61c6302416f
procutil: introduce context-manager interface for protect/restorestdio

The code looks slightly cleaner since it was pretty easy to pass arguments
in wrong order.

Patch

diff --git a/mercurial/commandserver.py b/mercurial/commandserver.py
--- a/mercurial/commandserver.py
+++ b/mercurial/commandserver.py
@@ -318,13 +318,12 @@  class pipeservice(object):
         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 = procutil.protectstdio(ui.fin, ui.fout)
-        try:
-            sv = server(ui, self.repo, fin, fout)
-            return sv.serve()
-        finally:
-            sv.cleanup()
-            procutil.restorestdio(ui.fin, ui.fout, fin, fout)
+        with procutil.protectedstdio(ui.fin, ui.fout) as (fin, fout):
+            try:
+                sv = server(ui, self.repo, fin, fout)
+                return sv.serve()
+            finally:
+                sv.cleanup()
 
 def _initworkerprocess():
     # use a different process group from the master process, in order to:
diff --git a/mercurial/utils/procutil.py b/mercurial/utils/procutil.py
--- a/mercurial/utils/procutil.py
+++ b/mercurial/utils/procutil.py
@@ -9,6 +9,7 @@ 
 
 from __future__ import absolute_import
 
+import contextlib
 import imp
 import io
 import os
@@ -240,6 +241,15 @@  def restorestdio(uin, uout, fin, fout):
             os.dup2(f.fileno(), uif.fileno())
             f.close()
 
+@contextlib.contextmanager
+def protectedstdio(uin, uout):
+    """Run code block with protected standard streams"""
+    fin, fout = protectstdio(uin, uout)
+    try:
+        yield fin, fout
+    finally:
+        restorestdio(uin, uout, fin, fout)
+
 def shellenviron(environ=None):
     """return environ with optional override, useful for shelling out"""
     def py2shell(val):