[PATCH 3 of 6] profiling: make profiling functions context managers (API)

[Gregory Szorc]

# HG changeset patch
# User Gregory Szorc <gregory.szorc at gmail.com>
# Date 1471217557 25200
#      Sun Aug 14 16:32:37 2016 -0700
# Node ID cf7b933cbc7fd0dfc4c5d5d67deae2d52866088a
# Parent  fd888ffaab6720688e4ad3f0358be09509effb6f
profiling: make profiling functions context managers (API)

This makes profiling more flexible since we can now call multiple
functions when a profiler is active. But the real reason for this
is to enable a future consumer to profile a function that returns
a generator. We can't do this from the profiling function itself
because functions can either be generators or have return values:
they can't be both. So therefore it isn't possible to have a generic
profiling function that can both consume and re-emit a generator
and return a value.

diff --git a/mercurial/dispatch.py b/mercurial/dispatch.py
--- a/mercurial/dispatch.py
+++ b/mercurial/dispatch.py
@@ -904,17 +904,18 @@ def _runcommand(ui, options, cmd, cmdfun
     """
     def checkargs():
         try:
             return cmdfunc()
         except error.SignatureError:
             raise error.CommandError(cmd, _("invalid arguments"))
 
     if ui.configbool('profiling', 'enabled'):
-        return profiling.profile(ui, checkargs)
+        with profiling.profile(ui):
+            return checkargs()
     else:
         return checkargs()
 
 def _exceptionwarning(ui):
     """Produce a warning message for the current active exception"""
 
     # For compatibility checking, we discard the portion of the hg
     # version after the + on the assumption that if a "normal
diff --git a/mercurial/profiling.py b/mercurial/profiling.py
--- a/mercurial/profiling.py
+++ b/mercurial/profiling.py
@@ -2,27 +2,29 @@
 #
 # Copyright 2016 Gregory Szorc <gregory.szorc at gmail.com>
 #
 # This software may be used and distributed according to the terms of the
 # GNU General Public License version 2 or any later version.
 
 from __future__ import absolute_import, print_function
 
+import contextlib
 import os
 import sys
 import time
 
 from .i18n import _
 from . import (
     error,
     util,
 )
 
-def lsprofile(ui, func, fp):
+ at contextlib.contextmanager
+def lsprofile(ui, fp):
     format = ui.config('profiling', 'format', default='text')
     field = ui.config('profiling', 'sort', default='inlinetime')
     limit = ui.configint('profiling', 'limit', default=30)
     climit = ui.configint('profiling', 'nested', default=0)
 
     if format not in ['text', 'kcachegrind']:
         ui.warn(_("unrecognized profiling format '%s'"
                     " - Ignored\n") % format)
@@ -32,76 +34,82 @@ def lsprofile(ui, func, fp):
         from . import lsprof
     except ImportError:
         raise error.Abort(_(
             'lsprof not available - install from '
             'http://codespeak.net/svn/user/arigo/hack/misc/lsprof/'))
     p = lsprof.Profiler()
     p.enable(subcalls=True)
     try:
-        return func()
+        yield
     finally:
         p.disable()
 
         if format == 'kcachegrind':
             from . import lsprofcalltree
             calltree = lsprofcalltree.KCacheGrind(p)
             calltree.output(fp)
         else:
             # format == 'text'
             stats = lsprof.Stats(p.getstats())
             stats.sort(field)
             stats.pprint(limit=limit, file=fp, climit=climit)
 
-def flameprofile(ui, func, fp):
+ at contextlib.contextmanager
+def flameprofile(ui, fp):
     try:
         from flamegraph import flamegraph
     except ImportError:
         raise error.Abort(_(
             'flamegraph not available - install from '
             'https://github.com/evanhempel/python-flamegraph'))
     # developer config: profiling.freq
     freq = ui.configint('profiling', 'freq', default=1000)
     filter_ = None
     collapse_recursion = True
     thread = flamegraph.ProfileThread(fp, 1.0 / freq,
                                       filter_, collapse_recursion)
     start_time = time.clock()
     try:
         thread.start()
-        func()
+        yield
     finally:
         thread.stop()
         thread.join()
         print('Collected %d stack frames (%d unique) in %2.2f seconds.' % (
             time.clock() - start_time, thread.num_frames(),
             thread.num_frames(unique=True)))
 
-def statprofile(ui, func, fp):
+def statprofile(ui, fp):
     try:
         import statprof
     except ImportError:
         raise error.Abort(_(
             'statprof not available - install using "easy_install statprof"'))
 
     freq = ui.configint('profiling', 'freq', default=1000)
     if freq > 0:
         statprof.reset(freq)
     else:
         ui.warn(_("invalid sampling frequency '%s' - ignoring\n") % freq)
 
     statprof.start()
     try:
-        return func()
+        yield
     finally:
         statprof.stop()
         statprof.display(fp)
 
-def profile(ui, fn):
-    """Profile a function call."""
+ at contextlib.contextmanager
+def profile(ui):
+    """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 = os.getenv('HGPROF')
     if profiler is None:
         profiler = ui.config('profiling', 'type', default='ls')
     if profiler not in ('ls', 'stat', 'flame'):
         ui.warn(_("unrecognized profiler '%s' - ignored\n") % profiler)
         profiler = 'ls'
 
     output = ui.config('profiling', 'output')
@@ -111,21 +119,25 @@ def profile(ui, fn):
     elif output:
         path = ui.expandpath(output)
         fp = open(path, 'wb')
     else:
         fp = sys.stderr
 
     try:
         if profiler == 'ls':
-            return lsprofile(ui, fn, fp)
+            proffn = lsprofile
         elif profiler == 'flame':
-            return flameprofile(ui, fn, fp)
+            proffn = flameprofile
         else:
-            return statprofile(ui, fn, fp)
+            proffn = statprofile
+
+        with proffn(ui, fp):
+            yield
+
     finally:
         if output:
             if output == 'blackbox':
                 val = 'Profile:\n%s' % 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)