diff: make use of output labeling

diff: make use of output labeling

--- a/mercurial/commands.py	Fri Apr 02 15:22:05 2010 -0500
+++ b/mercurial/commands.py	Fri Apr 02 15:22:06 2010 -0500
@@ -1174,14 +1174,16 @@
     diffopts = patch.diffopts(ui, opts)
 
     m = cmdutil.match(repo, pats, opts)
-    it = patch.diff(repo, node1, node2, match=m, opts=diffopts)
     if stat:
+        it = patch.diff(repo, node1, node2, match=m, opts=diffopts)
         width = ui.interactive() and util.termwidth() or 80
-        ui.write(patch.diffstat(util.iterlines(it), width=width,
-                                git=diffopts.git))
+        for chunk, label in patch.diffstatui(util.iterlines(it), width=width,
+                                             git=diffopts.git):
+            ui.write(chunk, label=label)
     else:
-        for chunk in it:
-            ui.write(chunk)
+        it = patch.diffui(repo, node1, node2, match=m, opts=diffopts)
+        for chunk, label in it:
+            ui.write(chunk, label=label)
 
 def export(ui, repo, *changesets, **opts):
     """dump the header and diffs for one or more changesets

--- a/mercurial/patch.py	Fri Apr 02 15:22:05 2010 -0500
+++ b/mercurial/patch.py	Fri Apr 02 15:22:06 2010 -0500
@@ -1466,6 +1466,43 @@
     else:
         return difffn(opts, None)
 
+def difflabel(func, *args, **kw):
+    '''yields 2-tuples of (output, label) based on the output of func()'''
+    prefixes = [('diff', 'diff.diffline'),
+                ('copy', 'diff.extended'),
+                ('rename', 'diff.extended'),
+                ('old', 'diff.extended'),
+                ('new', 'diff.extended'),
+                ('deleted', 'diff.extended'),
+                ('---', 'diff.file_a'),
+                ('+++', 'diff.file_b'),
+                ('@@', 'diff.hunk'),
+                ('-', 'diff.deleted'),
+                ('+', 'diff.inserted')]
+
+    for chunk in func(*args, **kw):
+        lines = chunk.split('\n')
+        for i, line in enumerate(lines):
+            if i != 0:
+                yield ('\n', '')
+            stripline = line
+            if line and line[0] in '+-':
+                # highlight trailing whitespace, but only in changed lines
+                stripline = line.rstrip()
+            for prefix, label in prefixes:
+                if stripline.startswith(prefix):
+                    yield (stripline, label)
+                    break
+            else:
+                yield (line, '')
+            if line != stripline:
+                yield (line[len(stripline):], 'diff.trailingwhitespace')
+
+def diffui(*args, **kw):
+    '''like diff(), but yields 2-tuples of (output, label) for ui.write()'''
+    return difflabel(diff, *args, **kw)
+
+
 def _addmodehdr(header, omode, nmode):
     if omode != nmode:
         header.append('old mode %s\n' % omode)
@@ -1636,3 +1673,22 @@
                       % (len(stats), totaladds, totalremoves))
 
     return ''.join(output)
+
+def diffstatui(*args, **kw):
+    '''like diffstat(), but yields 2-tuples of (output, label) for
+    ui.write()
+    '''
+
+    for line in diffstat(*args, **kw).splitlines():
+        if line and line[-1] in '+-':
+            name, graph = line.rsplit(' ', 1)
+            yield (name + ' ', '')
+            m = re.search(r'\++', graph)
+            if m:
+                yield (m.group(0), 'diffstat.inserted')
+            m = re.search(r'-+', graph)
+            if m:
+                yield (m.group(0), 'diffstat.deleted')
+        else:
+            yield (line, '')
+        yield ('\n', '')