# HG changeset patch
# User Gregory Szorc <gregory.szorc@gmail.com>
# Date 2016-08-15 01:25:22
# Node ID 5d44197c208b4e26ed2fa1eb627d5dea1ee50220
# Parent  97bfc2e5fba5f83981a5150d1401ecd50ee5f549

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
@@ -909,7 +909,8 @@ def _runcommand(ui, options, cmd, cmdfun
             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()
 
diff --git a/mercurial/profiling.py b/mercurial/profiling.py
--- a/mercurial/profiling.py
+++ b/mercurial/profiling.py
@@ -7,6 +7,7 @@
 
 from __future__ import absolute_import, print_function
 
+import contextlib
 import os
 import sys
 import time
@@ -17,7 +18,8 @@ from . import (
     util,
 )
 
-def lsprofile(ui, func, fp):
+@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)
@@ -37,7 +39,7 @@ def lsprofile(ui, func, fp):
     p = lsprof.Profiler()
     p.enable(subcalls=True)
     try:
-        return func()
+        yield
     finally:
         p.disable()
 
@@ -51,7 +53,8 @@ def lsprofile(ui, func, fp):
             stats.sort(field)
             stats.pprint(limit=limit, file=fp, climit=climit)
 
-def flameprofile(ui, func, fp):
+@contextlib.contextmanager
+def flameprofile(ui, fp):
     try:
         from flamegraph import flamegraph
     except ImportError:
@@ -67,7 +70,7 @@ def flameprofile(ui, func, fp):
     start_time = time.clock()
     try:
         thread.start()
-        func()
+        yield
     finally:
         thread.stop()
         thread.join()
@@ -75,7 +78,8 @@ def flameprofile(ui, func, fp):
             time.clock() - start_time, thread.num_frames(),
             thread.num_frames(unique=True)))
 
-def statprofile(ui, func, fp):
+@contextlib.contextmanager
+def statprofile(ui, fp):
     try:
         import statprof
     except ImportError:
@@ -90,13 +94,18 @@ def statprofile(ui, func, fp):
 
     statprof.start()
     try:
-        return func()
+        yield
     finally:
         statprof.stop()
         statprof.display(fp)
 
-def profile(ui, fn):
-    """Profile a function call."""
+@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')
@@ -116,11 +125,15 @@ def profile(ui, fn):
 
     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':