gh-133722: Add Difflib theme to _colorize and 'color' option to difflib.unified_diff (#133725)

Lib/difflib.py

@@ -30,6 +30,7 @@
            'Differ','IS_CHARACTER_JUNK', 'IS_LINE_JUNK', 'context_diff',
            'unified_diff', 'diff_bytes', 'HtmlDiff', 'Match']
 
+from _colorize import can_colorize, get_theme
 from heapq import nlargest as _nlargest
 from collections import namedtuple as _namedtuple
 from types import GenericAlias
@@ -1094,7 +1095,7 @@ def _format_range_unified(start, stop):
     return '{},{}'.format(beginning, length)
 
 def unified_diff(a, b, fromfile='', tofile='', fromfiledate='',
-                 tofiledate='', n=3, lineterm='\n'):
+                 tofiledate='', n=3, lineterm='\n', *, color=False):
     r"""
     Compare two sequences of lines; generate the delta as a unified diff.
 
@@ -1111,6 +1112,10 @@ def unified_diff(a, b, fromfile='', tofile='', fromfiledate='',
     For inputs that do not have trailing newlines, set the lineterm
     argument to "" so that the output will be uniformly newline free.
 
+    Set 'color' to True to enable output in color, similar to
+    'git diff --color'. Even if enabled, it can be
+    controlled using environment variables such as 'NO_COLOR'.
+
     The unidiff format normally has a header for filenames and modification
     times.  Any or all of these may be specified using strings for
     'fromfile', 'tofile', 'fromfiledate', and 'tofiledate'.
@@ -1134,6 +1139,11 @@ def unified_diff(a, b, fromfile='', tofile='', fromfiledate='',
      four
     """
 
+    if color and can_colorize():
+        t = get_theme(force_color=True).difflib
+    else:
+        t = get_theme(force_no_color=True).difflib
+
     _check_types(a, b, fromfile, tofile, fromfiledate, tofiledate, lineterm)
     started = False
     for group in SequenceMatcher(None,a,b).get_grouped_opcodes(n):
@@ -1141,25 +1151,25 @@ def unified_diff(a, b, fromfile='', tofile='', fromfiledate='',
             started = True
             fromdate = '\t{}'.format(fromfiledate) if fromfiledate else ''
             todate = '\t{}'.format(tofiledate) if tofiledate else ''
-            yield '--- {}{}{}'.format(fromfile, fromdate, lineterm)
-            yield '+++ {}{}{}'.format(tofile, todate, lineterm)
+            yield f'{t.header}--- {fromfile}{fromdate}{lineterm}{t.reset}'
+            yield f'{t.header}+++ {tofile}{todate}{lineterm}{t.reset}'
 
         first, last = group[0], group[-1]
         file1_range = _format_range_unified(first[1], last[2])
         file2_range = _format_range_unified(first[3], last[4])
-        yield '@@ -{} +{} @@{}'.format(file1_range, file2_range, lineterm)
+        yield f'{t.hunk}@@ -{file1_range} +{file2_range} @@{lineterm}{t.reset}'
 
         for tag, i1, i2, j1, j2 in group:
             if tag == 'equal':
                 for line in a[i1:i2]:
-                    yield ' ' + line
+                    yield f'{t.context} {line}{t.reset}'
                 continue
             if tag in {'replace', 'delete'}:
                 for line in a[i1:i2]:
-                    yield '-' + line
+                    yield f'{t.removed}-{line}{t.reset}'
             if tag in {'replace', 'insert'}:
                 for line in b[j1:j2]:
-                    yield '+' + line
+                    yield f'{t.added}+{line}{t.reset}'
 
 
 ########################################################################