[PATCH] color: add support for the 'export' command

[Dan Villiom Podlaski Christiansen]

# HG changeset patch
# User Dan Villiom Podlaski Christiansen <danchr at gmail.com>
# Date 1260000416 -3600
# Node ID cf5500afd34636e92e19620ebdadeffe6962f067
# Parent  b30efd0b2ced5bcd1ee61d2406337804ed47e59d
color: add support for the 'export' command.

diff --git a/hgext/color.py b/hgext/color.py
--- a/hgext/color.py
+++ b/hgext/color.py
@@ -223,6 +223,61 @@ def colorchurn(orig, ui, repo, *pats, **
     finally:
         ui.write = oldwrite
 
+def colorexport(orig, ui, repo, *pats, **opts):
+    '''run the export command with colored output
+
+       A stateful wrapper of repo.ui.write achieves this.'''
+    # Python doesn't do tail call optimisation, so all those nice
+    # tail calls make no difference...
+
+    def write(orig, s, newhandler=None):
+        '''Forward the call to the handler, possibly the one given'''
+        if newhandler:
+            handler[0] = newhandler
+        handler[0](orig, s)
+
+    def header(orig, s):
+        '''Write the header until a line that doesn't start with "# "'''
+        if s.startswith('# '):
+            orig(render_effects(s, _diff_effects['header']))
+        else:
+            write(orig, s, message)
+
+    def message(orig, s):
+        '''
+        If the string starts with a diff line, it's a diff. Messages
+        looking like may not be handled gracefully.
+        '''
+        if s.startswith('diff --git a/') or s.startswith('diff a/'):
+            write(orig, s, diff)
+        else:
+            # the effects may not span multiple lines
+            for l in s.splitlines(True):
+                orig(render_effects(l, _diff_effects['message']))
+
+    def diff(orig, s):
+        '''
+        Write the diff/patch until a line starting with '# ' is seen.
+        This means the patch is done and the header of the next
+        follows.
+        '''
+        if s and s.startswith('# '):
+            write(orig, s, header)
+        else:
+            colorwrap(orig, s)
+
+    # use a list for state so the nested functions can access and modify it
+    handler = [header]
+
+    # do the actual wrapping
+    oldwrite = extensions.wrapfunction(repo.ui, 'write', write)
+
+    try:
+        orig(ui, repo, *pats, **opts)
+    finally:
+        repo.ui.write = oldwrite
+
+
 _diff_prefixes = [('diff', 'diffline'),
                   ('copy', 'extended'),
                   ('rename', 'extended'),
@@ -240,6 +295,8 @@ _diff_effects = {'diffline': ['bold'],
                  'file_a': ['red', 'bold'],
                  'file_b': ['green', 'bold'],
                  'hunk': ['magenta'],
+                 'header': ['yellow', 'bold'],
+                 'message': ['blue'],
                  'deleted': ['red'],
                  'inserted': ['green'],
                  'changed': ['white'],
@@ -247,6 +304,7 @@ _diff_effects = {'diffline': ['bold'],
 
 def extsetup(ui):
     '''Initialize the extension.'''
+    _setupcmd(ui, 'export', commands.table, colorexport, _diff_effects)
     _setupcmd(ui, 'diff', commands.table, colordiff, _diff_effects)
     _setupcmd(ui, 'incoming', commands.table, None, _diff_effects)
     _setupcmd(ui, 'log', commands.table, None, _diff_effects)
diff --git a/mercurial/patch.py b/mercurial/patch.py
--- a/mercurial/patch.py
+++ b/mercurial/patch.py
@@ -1340,7 +1340,13 @@ def export(repo, revs, template='hg-%h.p
             fp = cmdutil.make_file(repo, template, node, total=total,
                                    seqno=seqno, revwidth=revwidth,
                                    mode='ab')
-        if fp != sys.stdout and hasattr(fp, 'name'):
+
+        if fp == sys.stdout:
+            # The color extension relies on UI writes being routed to the
+            # UI instance, and the standard UI class is hard coded to write to
+            # sys.stdout & sys.stderr for output. So we hard code it too.
+            fp = repo.ui
+        elif hasattr(fp, 'name'):
             repo.ui.note("%s\n" % fp.name)
 
         fp.write("# HG changeset patch\n")