mercurial/profiling.py
changeset 32783 4483696dacee
parent 32417 f40dc6f7c12f
child 32784 086c1ef0f666
--- a/mercurial/profiling.py	Fri Jun 09 22:15:53 2017 -0400
+++ b/mercurial/profiling.py	Thu Jun 08 01:38:48 2017 +0100
@@ -141,35 +141,41 @@
 
         statprof.display(fp, data=data, format=displayformat, **kwargs)
 
-@contextlib.contextmanager
-def profile(ui):
+class profile(object):
     """Start profiling.
 
     Profiling is active when the context manager is active. When the context
     manager exits, profiling results will be written to the configured output.
     """
-    profiler = encoding.environ.get('HGPROF')
-    proffn = None
-    if profiler is None:
-        profiler = ui.config('profiling', 'type', default='stat')
-    if profiler not in ('ls', 'stat', 'flame'):
-        # try load profiler from extension with the same name
-        proffn = _loadprofiler(ui, profiler)
-        if proffn is None:
-            ui.warn(_("unrecognized profiler '%s' - ignored\n") % profiler)
-            profiler = 'stat'
+    def __init__(self, ui):
+        self._ui = ui
+        self._output = None
+        self._fp = None
+        self._profiler = None
 
-    output = ui.config('profiling', 'output')
+    def __enter__(self):
+        profiler = encoding.environ.get('HGPROF')
+        proffn = None
+        if profiler is None:
+            profiler = self._ui.config('profiling', 'type', default='stat')
+        if profiler not in ('ls', 'stat', 'flame'):
+            # try load profiler from extension with the same name
+            proffn = _loadprofiler(self._ui, profiler)
+            if proffn is None:
+                self._ui.warn(_("unrecognized profiler '%s' - ignored\n")
+                              % profiler)
+                profiler = 'stat'
 
-    if output == 'blackbox':
-        fp = util.stringio()
-    elif output:
-        path = ui.expandpath(output)
-        fp = open(path, 'wb')
-    else:
-        fp = ui.ferr
+        self._output = self._ui.config('profiling', 'output')
 
-    try:
+        if self._output == 'blackbox':
+            self._fp = util.stringio()
+        elif self._output:
+            path = self._ui.expandpath(self._output)
+            self._fp = open(path, 'wb')
+        else:
+            self._fp = self._ui.ferr
+
         if proffn is not None:
             pass
         elif profiler == 'ls':
@@ -179,18 +185,21 @@
         else:
             proffn = statprofile
 
-        with proffn(ui, fp):
-            yield
+        self._profiler = proffn(self._ui, self._fp)
+        self._profiler.__enter__()
 
-    finally:
-        if output:
-            if output == 'blackbox':
-                val = 'Profile:\n%s' % fp.getvalue()
+    def __exit__(self, exception_type, exception_value, traceback):
+        if self._profiler is None:
+            return
+        self._profiler.__exit__(exception_type, exception_value, traceback)
+        if self._output:
+            if self._output == 'blackbox':
+                val = 'Profile:\n%s' % self._fp.getvalue()
                 # ui.log treats the input as a format string,
                 # so we need to escape any % signs.
                 val = val.replace('%', '%%')
-                ui.log('profile', val)
-            fp.close()
+                self._ui.log('profile', val)
+            self._fp.close()
 
 @contextlib.contextmanager
 def maybeprofile(ui):