[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 1248721832 -7200
# Node ID 3fa9f0e6dd95eb48fb24d76442402e6ed81903e3
# Parent  83aa93f22680731176dea2f95ce1e8cace538363
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
@@ -197,6 +197,68 @@ def colordiff(orig, ui, repo, *pats, **o
     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, diff_or_message)
+
+    def diff_or_message(orig, s):
+        '''
+        if the string starts with a diff line, it's a diff.
+        Messages looking like diffs aren't handled gracefully.
+        '''
+        if s.startswith('diff --git a/') or s.startswith('diff a/'):
+            write(orig, s, diff)
+        else:
+            write(orig, s, 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)
+
+    def message(orig, s):
+        '''
+        Write the commit message. A write consisting of two newlines
+        following eachother may mean the message is written and the
+        patch follows.
+        '''
+        orig(render_effects(s, _diff_effects['message']))
+        if s == '\n\n':
+            handler[0] = diff_or_message
+
+    # 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'),
@@ -214,6 +276,8 @@ _diff_effects = {'diffline': ['bold'],
                  'file_a': ['red', 'bold'],
                  'file_b': ['green', 'bold'],
                  'hunk': ['magenta'],
+                 'header': ['cyan', 'italic'],
+                 'message': ['blue', 'bold'],
                  'deleted': ['red'],
                  'inserted': ['green'],
                  'changed': ['white'],
@@ -225,6 +289,7 @@ def uisetup(ui):
     '''Initialize the extension.'''
     global _ui
     _ui = ui
+    _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
@@ -1364,7 +1364,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")