# HG changeset patch
# User Jun Wu <quark@fb.com>
# Date 2017-05-22 08:17:49
# Node ID f40dc6f7c12f36fb56dbadeb32a96ee25e6bacbb
# Parent  9a3e88d4a0300da375edebaf9f43fffddd5568cf

profiling: allow loading profiling extension before everything else

6d642ecf1a89 makes profiler start early without loading extensions. That
makes it impossible for an extension to add customized profilers.

This patch adds a special case: if a profiler is not found but an extension
with the same name could be loaded, load that extension first, and expect it
to have a "profile" contextmanager method. This allows customized profilers
and extension setup time is still profiled.

diff --git a/mercurial/extensions.py b/mercurial/extensions.py
--- a/mercurial/extensions.py
+++ b/mercurial/extensions.py
@@ -181,7 +181,7 @@ def _runextsetup(name, ui):
 
 def loadall(ui, whitelist=None):
     result = ui.configitems("extensions")
-    if whitelist:
+    if whitelist is not None:
         result = [(k, v) for (k, v) in result if k in whitelist]
     newindex = len(_order)
     for (name, path) in result:
diff --git a/mercurial/profiling.py b/mercurial/profiling.py
--- a/mercurial/profiling.py
+++ b/mercurial/profiling.py
@@ -13,9 +13,21 @@ from .i18n import _
 from . import (
     encoding,
     error,
+    extensions,
     util,
 )
 
+def _loadprofiler(ui, profiler):
+    """load profiler extension. return profile method, or None on failure"""
+    extname = profiler
+    extensions.loadall(ui, whitelist=[extname])
+    try:
+        mod = extensions.find(extname)
+    except KeyError:
+        return None
+    else:
+        return getattr(mod, 'profile', None)
+
 @contextlib.contextmanager
 def lsprofile(ui, fp):
     format = ui.config('profiling', 'format', default='text')
@@ -137,11 +149,15 @@ def profile(ui):
     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'):
-        ui.warn(_("unrecognized profiler '%s' - ignored\n") % profiler)
-        profiler = 'stat'
+        # 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'
 
     output = ui.config('profiling', 'output')
 
@@ -154,7 +170,9 @@ def profile(ui):
         fp = ui.ferr
 
     try:
-        if profiler == 'ls':
+        if proffn is not None:
+            pass
+        elif profiler == 'ls':
             proffn = lsprofile
         elif profiler == 'flame':
             proffn = flameprofile
diff --git a/tests/test-profile.t b/tests/test-profile.t
--- a/tests/test-profile.t
+++ b/tests/test-profile.t
@@ -99,3 +99,51 @@ statprof can be used as a standalone mod
   [1]
 
   $ cd ..
+
+profiler extension could be loaded before other extensions
+
+  $ cat > fooprof.py <<EOF
+  > from __future__ import absolute_import
+  > import contextlib
+  > @contextlib.contextmanager
+  > def profile(ui, fp):
+  >     print('fooprof: start profile')
+  >     yield
+  >     print('fooprof: end profile')
+  > def extsetup(ui):
+  >     ui.write('fooprof: loaded\n')
+  > EOF
+
+  $ cat > otherextension.py <<EOF
+  > from __future__ import absolute_import
+  > def extsetup(ui):
+  >     ui.write('otherextension: loaded\n')
+  > EOF
+
+  $ hg init b
+  $ cd b
+  $ cat >> .hg/hgrc <<EOF
+  > [extensions]
+  > other = $TESTTMP/otherextension.py
+  > fooprof = $TESTTMP/fooprof.py
+  > EOF
+
+  $ hg root
+  otherextension: loaded
+  fooprof: loaded
+  $TESTTMP/b (glob)
+  $ HGPROF=fooprof hg root --profile
+  fooprof: loaded
+  fooprof: start profile
+  otherextension: loaded
+  $TESTTMP/b (glob)
+  fooprof: end profile
+
+  $ HGPROF=other hg root --profile 2>&1 | head -n 2
+  otherextension: loaded
+  unrecognized profiler 'other' - ignored
+
+  $ HGPROF=unknown hg root --profile 2>&1 | head -n 1
+  unrecognized profiler 'unknown' - ignored
+
+  $ cd ..