More work on the diff kitten

kittens/diff/collect.py

@@ -91,7 +91,7 @@ def collect_files(collection, left, right):
         collection.add_add(right_path_map[name])
 
 
-sanitize_pat = re.compile('[\x00-\x1f\x7f\x80-\x9f]')
+sanitize_pat = re.compile('[\x00-\x09\x0b-\x1f\x7f\x80-\x9f]')
 
 
 def sanitize_sub(m):
@@ -127,7 +127,7 @@ def data_for_path(path):
 @lru_cache(maxsize=1024)
 def lines_for_path(path):
     data = data_for_path(path)
-    return tuple(map(sanitize, data.splitlines()))
+    return tuple(sanitize(data).splitlines())
 
 
 @lru_cache(maxsize=1024)
@@ -146,3 +146,15 @@ def create_collection(left, right):
         collection.add_change(pl, pr)
     collection.finalize()
     return collection
+
+
+highlight_data = {}
+
+
+def set_highlight_data(data):
+    global highlight_data
+    highlight_data = data
+
+
+def highlights_for_path(path):
+    return highlight_data.get(path, ())

kittens/diff/highlight.py

@@ -2,6 +2,10 @@
 # vim:fileencoding=utf-8
 # License: GPL v3 Copyright: 2018, Kovid Goyal <kovid at kovidgoyal.net>
 
+import concurrent
+import os
+import re
+
 from pygments import highlight
 from pygments.formatter import Formatter
 from pygments.lexers import get_lexer_for_filename
@@ -9,6 +13,8 @@ from pygments.util import ClassNotFound
 
 from kitty.rgb import color_as_sgr, parse_sharp
 
+from .collect import data_for_path, lines_for_path
+
 
 class DiffFormatter(Formatter):
 
@@ -76,6 +82,64 @@ def highlight_data(code, filename):
         return highlight(code, lexer, formatter)
 
 
+split_pat = re.compile(r'(\033\[.*?m)')
+
+
+class Segment:
+
+    __slots__ = ('start', 'end', 'start_code', 'end_code')
+
+    def __init__(self, start, start_code):
+        self.start = start
+        self.start_code = start_code
+
+
+def highlight_line(line):
+    ans = []
+    current = None
+    pos = 0
+    for x in split_pat.split(line):
+        if x.startswith('\033'):
+            if current is None:
+                current = Segment(pos, x)
+            else:
+                current.end = pos
+                current.end_code = x
+                ans.append(current)
+                current = None
+        else:
+            pos += len(x)
+    return ans
+
+
+def highlight_for_diff(path):
+    ans = []
+    lines = lines_for_path(path)
+    hd = highlight_data('\n'.join(lines), path)
+    if hd is not None:
+        for line in hd.splitlines():
+            ans.append(highlight_line(line))
+    return ans
+
+
+def highlight_collection(collection):
+    jobs = {}
+    ans = {}
+    with concurrent.futures.ProcessPoolExecutor(max_workers=os.cpu_count()) as executor:
+        for path, item_type, other_path in collection:
+            is_binary = isinstance(data_for_path(path), bytes)
+            if not is_binary:
+                jobs[executor.submit(highlight_for_diff, path)] = path
+        for future in concurrent.futures.as_completed(jobs):
+            path = jobs[future]
+            try:
+                highlights = future.result()
+            except Exception as e:
+                return 'Running syntax highlighting for {} generated an exception: {}'.format(path, e)
+            ans[path] = highlights
+    return ans
+
+
 def main():
     # kitty +runpy "from kittens.diff.highlight import main; main()" file
     import sys

kittens/diff/main.py

@@ -14,11 +14,17 @@ from kitty.key_encoding import (
 
 from ..tui.handler import Handler
 from ..tui.loop import Loop
-from .collect import create_collection, data_for_path
+from .collect import create_collection, data_for_path, set_highlight_data
 from .config import init_config
 from .patch import Differ
 from .render import render_diff
 
+try:
+    from .highlight import initialize_highlighter, highlight_collection
+except ImportError:
+    initialize_highlighter = None
+
+
 INITIALIZING, COLLECTED, DIFFED = range(3)
 
 
@@ -180,6 +186,18 @@ class DiffHandler(Handler):
             self.diff_map = diff_map
             self.render_diff()
             self.draw_screen()
+            if initialize_highlighter is not None:
+                initialize_highlighter()
+                self.start_job('highlight', highlight_collection, self.collection)
+        elif job_id == 'highlight':
+            hdata = job_result['result']
+            if isinstance(hdata, str):
+                self.report_traceback_on_exit = diff_map
+                self.quit_loop(1)
+                return
+            set_highlight_data(hdata)
+            self.render_diff()
+            self.draw_screen()
 
     def on_interrupt(self):
         self.quit_loop(1)

kittens/diff/render.py

@@ -7,7 +7,10 @@ from itertools import repeat
 
 from kitty.fast_data_types import truncate_point_for_length, wcswidth
 
-from .collect import data_for_path, lines_for_path, path_name_map, sanitize
+from .collect import (
+    data_for_path, highlights_for_path, lines_for_path, path_name_map,
+    sanitize
+)
 from .config import formats
 
 
@@ -198,6 +201,18 @@ class DiffData:
         self.filler_line = render_diff_line('', '', 'filler', margin_size, available_cols)
         self.left_filler_line = render_diff_line('', '', 'remove', margin_size, available_cols)
         self.right_filler_line = render_diff_line('', '', 'add', margin_size, available_cols)
+        self.left_hdata = highlights_for_path(left_path)
+        self.right_hdata = highlights_for_path(right_path)
+
+    def left_highlights_for_line(self, line_num):
+        if line_num < len(self.left_hdata):
+            return self.left_hdata[line_num]
+        return ()
+
+    def right_highlights_for_line(self, line_num):
+        if line_num < len(self.right_hdata):
+            return self.right_hdata[line_num]
+        return ()
 
 
 def render_diff_line(number, text, ltype, margin_size, available_cols):
@@ -221,12 +236,12 @@ def hunk_title(hunk_num, hunk, margin_size, available_cols):
     return m + hunk_format(place_in(t, available_cols))
 
 
-def render_half_line(line_number, src, ltype, margin_size, available_cols, changed_center=None):
+def render_half_line(line_number, line, highlights, ltype, margin_size, available_cols, changed_center=None):
     if changed_center is not None and changed_center[0]:
         start, stop = highlight_boundaries(ltype)
-        lines = split_to_size_with_center(src[line_number], available_cols, changed_center[0], changed_center[1], start, stop)
+        lines = split_to_size_with_center(line, available_cols, changed_center[0], changed_center[1], start, stop)
     else:
-        lines = split_to_size(src[line_number], available_cols)
+        lines = split_to_size(line, available_cols)
     line_number = str(line_number + 1)
     for line in lines:
         yield render_diff_line(line_number, line, ltype, margin_size, available_cols)
@@ -256,13 +271,17 @@ def lines_for_chunk(data, hunk_num, chunk, chunk_num):
             ref = Reference(data.left_path, HunkRef(hunk_num, chunk_num, i))
             ll, rl = [], []
             if i < chunk.left_count:
+                rln = chunk.left_start + i
                 ll.extend(render_half_line(
-                    chunk.left_start + i, data.left_lines, 'remove', data.margin_size,
+                    rln, data.left_lines[rln], data.left_highlights_for_line(rln),
-                    data.available_cols, None if chunk.centers is None else chunk.centers[i]))
+                    'remove', data.margin_size, data.available_cols,
+                    None if chunk.centers is None else chunk.centers[i]))
             if i < chunk.right_count:
+                rln = chunk.left_start + i
                 rl.extend(render_half_line(
-                    chunk.right_start + i, data.right_lines, 'add', data.margin_size,
+                    rln, data.right_lines[rln], data.right_highlights_for_line(rln),
-                    data.available_cols, None if chunk.centers is None else chunk.centers[i]))
+                    'add', data.margin_size, data.available_cols,
+                    None if chunk.centers is None else chunk.centers[i]))
             if i < common:
                 extra = len(ll) - len(rl)
                 if extra != 0:
@@ -298,8 +317,13 @@ def all_lines(path, args, columns, margin_size, is_add=True):
     lines = lines_for_path(path)
     filler = render_diff_line('', '', 'filler', margin_size, available_cols)
     msg_written = False
-    for line_number in range(len(lines)):
+    hdata = highlights_for_path(path)
-        h = render_half_line(line_number, lines, ltype, margin_size, available_cols)
+
+    def highlights(num):
+        return hdata[num] if num < len(hdata) else ()
+
+    for line_number, line in enumerate(lines):
+        h = render_half_line(line_number, line, highlights(line_number), ltype, margin_size, available_cols)
         for i, hl in enumerate(h):
             ref = Reference(path, LineRef(line_number, i))
             empty = filler