color/progress: subclass ui instead of using wrapfunction (issue2096)

This resolves the issue of hg cmd --mq not being colorized. This was due
to color wrapping only the instance of ui passed to dispatch._runcommand(),
which isn't the same ui object that mq.mqcommand() receives. After dispatch
calls extensions.loadall(), it makes sure any changes to ui.__class__ in
uisetup are propagated.

progress is updated to wrap ui in the same manner because wrapfunction
doesn't play well when ui.__class__ has been replaced by another extension
(orig will point to the old class method instead of color's).

hgext/color.py

@@ -74,7 +74,7 b" Any value other than 'ansi', 'win32', or"
 
 import os, sys
 
-from mercurial import commands, dispatch, extensions
+from mercurial import commands, dispatch, extensions, ui as uimod
 from mercurial.i18n import _
 
 # start and stop parameters for effects
@@ -140,49 +140,50 b' def configstyles(ui):'
                             % (e, status))
             _styles[status] = ' '.join(good)
 
-_buffers = None
-def style(msg, label):
-    effects = []
-    for l in label.split():
-        s = _styles.get(l, '')
-        if s:
-            effects.append(s)
-    effects = ''.join(effects)
-    if effects:
-        return '\n'.join([render_effects(s, effects)
-                          for s in msg.split('\n')])
-    return msg
+class colorui(uimod.ui):
+    def popbuffer(self, labeled=False):
+        if labeled:
+            return ''.join(self.label(a, label) for a, label
+                           in self._buffers.pop())
+        return ''.join(a for a, label in self._buffers.pop())
 
-def popbuffer(orig, labeled=False):
-    global _buffers
-    if labeled:
-        return ''.join(style(a, label) for a, label in _buffers.pop())
-    return ''.join(a for a, label in _buffers.pop())
+    _colormode = 'ansi'
+    def write(self, *args, **opts):
+        label = opts.get('label', '')
+        if self._buffers:
+            self._buffers[-1].extend([(str(a), label) for a in args])
+        elif self._colormode == 'win32':
+            for a in args:
+                win32print(a, orig, **opts)
+        else:
+            return super(colorui, self).write(
+                *[self.label(str(a), label) for a in args], **opts)
 
-mode = 'ansi'
-def write(orig, *args, **opts):
-    label = opts.get('label', '')
-    global _buffers
-    if _buffers:
-        _buffers[-1].extend([(str(a), label) for a in args])
-    elif mode == 'win32':
-        for a in args:
-            win32print(a, orig, **opts)
-    else:
-        return orig(*[style(str(a), label) for a in args], **opts)
+    def write_err(self, *args, **opts):
+        label = opts.get('label', '')
+        if self._colormode == 'win32':
+            for a in args:
+                win32print(a, orig, **opts)
+        else:
+            return super(colorui, self).write(
+                *[self.label(str(a), label) for a in args], **opts)
 
-def write_err(orig, *args, **opts):
-    label = opts.get('label', '')
-    if mode == 'win32':
-        for a in args:
-            win32print(a, orig, **opts)
-    else:
-        return orig(*[style(str(a), label) for a in args], **opts)
+    def label(self, msg, label):
+        effects = []
+        for l in label.split():
+            s = _styles.get(l, '')
+            if s:
+                effects.append(s)
+        effects = ''.join(effects)
+        if effects:
+            return '\n'.join([render_effects(s, effects)
+                              for s in msg.split('\n')])
+        return msg
+
 
 def uisetup(ui):
     if ui.plain():
         return
-    global mode
     mode = ui.config('color', 'mode', 'auto')
     if mode == 'auto':
         if os.name == 'nt' and 'TERM' not in os.environ:
@@ -202,14 +203,11 b' def uisetup(ui):'
         if (opts['color'] == 'always' or
             (opts['color'] == 'auto' and (os.environ.get('TERM') != 'dumb'
                                           and ui_.formatted()))):
-            global _buffers
-            _buffers = ui_._buffers
-            extensions.wrapfunction(ui_, 'popbuffer', popbuffer)
-            extensions.wrapfunction(ui_, 'write', write)
-            extensions.wrapfunction(ui_, 'write_err', write_err)
-            ui_.label = style
+            colorui._colormode = mode
+            colorui.__bases__ = (ui_.__class__,)
+            ui_.__class__ = colorui
             extstyles()
-            configstyles(ui)
+            configstyles(ui_)
         return orig(ui_, opts, cmd, cmdfunc)
     extensions.wrapfunction(dispatch, '_runcommand', colorcmd)
 

hgext/progress.py

@@ -45,7 +45,6 b' num characters, or ``+<num>`` for the fi'
 import sys
 import time
 
-from mercurial import extensions
 from mercurial import util
 
 def spacejoin(*args):
@@ -159,7 +158,7 b' class progbar(object):'
         tw = util.termwidth()
         return min(int(self.ui.config('progress', 'width', default=tw)), tw)
 
-    def progress(self, orig, topic, pos, item='', unit='', total=None):
+    def progress(self, topic, pos, item='', unit='', total=None):
         if pos is None:
             if self.topics and self.topics[-1] == topic and self.printed:
                 self.complete()
@@ -172,29 +171,35 b' class progbar(object):'
                 and topic == self.topics[-1]):
                 self.lastprint = now
                 self.show(topic, pos, item, unit, total)
-        return orig(topic, pos, item=item, unit=unit, total=total)
-
-    def write(self, orig, *args, **opts):
-        if self.printed:
-            self.clear()
-        return orig(*args, **opts)
-
-sharedprog = None
 
 def uisetup(ui):
+    class progressui(ui.__class__):
+        _progbar = None
+
+        def progress(self, *args, **opts):
+            self._progbar.progress(*args, **opts)
+            return super(progressui, self).progress(*args, **opts)
+
+        def write(self, *args, **opts):
+            if self._progbar.printed:
+                self._progbar.clear()
+            return super(progressui, self).write(*args, **opts)
+
+        def write_err(self, *args, **opts):
+            if self._progbar.printed:
+                self._progbar.clear()
+            return super(progressui, self).write_err(*args, **opts)
+
     # Apps that derive a class from ui.ui() can use
     # setconfig('progress', 'disable', 'True') to disable this extension
     if ui.configbool('progress', 'disable'):
         return
     if shouldprint(ui) and not ui.debugflag and not ui.quiet:
+        ui.__class__ = progressui
         # we instantiate one globally shared progress bar to avoid
         # competing progress bars when multiple UI objects get created
-        global sharedprog
-        if not sharedprog:
-            sharedprog = progbar(ui)
-        extensions.wrapfunction(ui, 'progress', sharedprog.progress)
-        extensions.wrapfunction(ui, 'write', sharedprog.write)
-        extensions.wrapfunction(ui, 'write_err', sharedprog.write)
+        if not progressui._progbar:
+            progressui._progbar = progbar(ui)
 
 def reposetup(ui, repo):
     uisetup(repo.ui)

mercurial/dispatch.py

@@ -388,6 +388,8 b' def _dispatch(ui, args):'
     # times so we keep track of configured extensions in _loaded.
     extensions.loadall(lui)
     exts = [ext for ext in extensions.extensions() if ext[0] not in _loaded]
+    # Propagate any changes to lui.__class__ by extensions
+    ui.__class__ = lui.__class__
 
     # (uisetup and extsetup are handled in extensions.loadall)
 

tests/test-mq

@@ -80,6 +80,9 b" echo ' .hgignore:'"
 cat .hg/patches/.hgignore
 echo '  series:'
 cat .hg/patches/series
+
+echo '% status --mq with color (issue2096)'
+hg status --mq --config extensions.color= --color=always
 cd ..
 
 echo '% init --mq without repo'

tests/test-mq.out

@@ -93,6 +93,11 b' bleh'
   series:
 A
 B
+% status --mq with color (issue2096)
+A .hgignore
+A A
+A B
+A series
 % init --mq without repo
 abort: There is no Mercurial repository here (.hg not found)
 % init --mq with repo path