Fix highlighting edges and unify highlighting code

babi/_types.py

@@ -2,7 +2,5 @@ from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from typing import Protocol  # python3.8+
-    from typing_extensions import TypedDict  # python3.8+
 else:
     Protocol = object
-    TypedDict = dict

babi/file.py

@@ -24,6 +24,7 @@ from typing import Union
 from babi.hl.interface import FileHL
 from babi.hl.interface import HLFactory
 from babi.hl.replace import Replace
+from babi.hl.selection import Selection
 from babi.horizontal_scrolling import line_x
 from babi.horizontal_scrolling import scrolled_line
 from babi.list_spy import ListSpy
@@ -139,7 +140,7 @@ def clear_selection(func: TCallable) -> TCallable:
     @functools.wraps(func)
     def clear_selection_inner(self: 'File', *args: Any, **kwargs: Any) -> Any:
         ret = func(self, *args, **kwargs)
-        self.select_start = None
+        self.selection.clear()
         return ret
     return cast(TCallable, clear_selection_inner)
 
@@ -220,9 +221,9 @@ class File:
         self.sha256: Optional[str] = None
         self.undo_stack: List[Action] = []
         self.redo_stack: List[Action] = []
-        self.select_start: Optional[Tuple[int, int]] = None
         self._hl_factories = hl_factories
         self._replace_hl = Replace()
+        self.selection = Selection()
         self._file_hls: Tuple[FileHL, ...] = ()
 
     def ensure_loaded(self, status: Status) -> None:
@@ -253,7 +254,7 @@ class File:
                 file_hls.append(factory.get_file_highlighter(self.filename))
             else:
                 file_hls.append(factory.get_blank_file_highlighter())
-        self._file_hls = (*file_hls, self._replace_hl)
+        self._file_hls = (*file_hls, self._replace_hl, self.selection)
 
     def __repr__(self) -> str:
         return f'<{type(self).__name__} {self.filename!r}>'
@@ -444,7 +445,7 @@ class File:
             self.x = self.x_hint = match.start()
             self.scroll_screen_if_needed(screen.margin)
             if res != 'a':  # make `a` replace the rest of them
-                with self._replace_hl.region(self.y, self.x, len(match[0])):
+                with self._replace_hl.region(self.y, self.x, match.end()):
                     screen.draw()
                     res = screen.quick_prompt(
                         'replace [y(es), n(o), a(ll)]?', 'yna',
@@ -537,16 +538,17 @@ class File:
 
     @edit_action('indent selection', final=True)
     def _indent_selection(self, margin: Margin) -> None:
-        assert self.select_start is not None
-        sel_y, sel_x = self.select_start
-        (s_y, _), (e_y, _) = self._get_selection()
+        assert self.selection.start is not None
+        sel_y, sel_x = self.selection.start
+        (s_y, _), (e_y, _) = self.selection.get()
         for l_y in range(s_y, e_y + 1):
             if self.lines[l_y]:
                 self.lines[l_y] = ' ' * 4 + self.lines[l_y]
-                if l_y == sel_y and sel_x != 0:
-                    self.select_start = (sel_y, sel_x + 4)
                 if l_y == self.y:
                     self.x = self.x_hint = self.x + 4
+                if l_y == sel_y and sel_x != 0:
+                    sel_x += 4
+        self.selection.set(sel_y, sel_x, self.y, self.x)
 
     @edit_action('insert tab', final=False)
     def _tab(self, margin: Margin) -> None:
@@ -557,7 +559,7 @@ class File:
         _restore_lines_eof_invariant(self.lines)
 
     def tab(self, margin: Margin) -> None:
-        if self.select_start:
+        if self.selection.start is not None:
             self._indent_selection(margin)
         else:
             self._tab(margin)
@@ -572,17 +574,18 @@ class File:
 
     @edit_action('dedent selection', final=True)
     def _dedent_selection(self, margin: Margin) -> None:
-        assert self.select_start is not None
-        sel_y, sel_x = self.select_start
-        (s_y, _), (e_y, _) = self._get_selection()
+        assert self.selection.start is not None
+        sel_y, sel_x = self.selection.start
+        (s_y, _), (e_y, _) = self.selection.get()
         for l_y in range(s_y, e_y + 1):
             n = self._dedent_line(self.lines[l_y])
             if n:
                 self.lines[l_y] = self.lines[l_y][n:]
-                if l_y == sel_y:
-                    self.select_start = (sel_y, max(sel_x - n, 0))
                 if l_y == self.y:
                     self.x = self.x_hint = max(self.x - n, 0)
+                if l_y == sel_y:
+                    sel_x = max(sel_x - n, 0)
+        self.selection.set(sel_y, sel_x, self.y, self.x)
 
     @edit_action('dedent', final=True)
     def _dedent(self, margin: Margin) -> None:
@@ -592,7 +595,7 @@ class File:
             self.x = self.x_hint = max(self.x - n, 0)
 
     def shift_tab(self, margin: Margin) -> None:
-        if self.select_start:
+        if self.selection.start is not None:
             self._dedent_selection(margin)
         else:
             self._dedent(margin)
@@ -601,7 +604,7 @@ class File:
     @clear_selection
     def cut_selection(self, margin: Margin) -> Tuple[str, ...]:
         ret = []
-        (s_y, s_x), (e_y, e_x) = self._get_selection()
+        (s_y, s_x), (e_y, e_x) = self.selection.get()
         if s_y == e_y:
             ret.append(self.lines[s_y][s_x:e_x])
             self.lines[s_y] = self.lines[s_y][:s_x] + self.lines[s_y][e_x:]
@@ -674,7 +677,7 @@ class File:
     @edit_action('sort selection', final=True)
     @clear_selection
     def sort_selection(self, margin: Margin) -> None:
-        (s_y, _), (e_y, _) = self._get_selection()
+        (s_y, _), (e_y, _) = self.selection.get()
         e_y = min(e_y + 1, len(self.lines) - 1)
         if self.lines[e_y - 1] == '':
             e_y -= 1
@@ -732,7 +735,7 @@ class File:
 
     def finalize_previous_action(self) -> None:
         assert not isinstance(self.lines, ListSpy), 'nested edit/movement'
-        self.select_start = None
+        self.selection.clear()
         if self.undo_stack:
             self.undo_stack[-1].final = True
 
@@ -785,14 +788,14 @@ class File:
 
     @contextlib.contextmanager
     def select(self) -> Generator[None, None, None]:
-        if self.select_start is None:
-            select_start = (self.y, self.x)
+        if self.selection.start is None:
+            start = (self.y, self.x)
         else:
-            select_start = self.select_start
+            start = self.selection.start
         try:
             yield
         finally:
-            self.select_start = select_start
+            self.selection.set(*start, self.y, self.x)
 
     # positioning
 
@@ -809,95 +812,50 @@ class File:
     ) -> None:
         stdscr.move(self.rendered_y(margin), self.rendered_x())
 
-    def _get_selection(self) -> Tuple[Tuple[int, int], Tuple[int, int]]:
-        assert self.select_start is not None
-        select_end = (self.y, self.x)
-        if select_end < self.select_start:
-            return select_end, self.select_start
-        else:
-            return self.select_start, select_end
-
     def touch(self, lineno: int) -> None:
         for file_hl in self._file_hls:
             file_hl.touch(lineno)
 
     def draw(self, stdscr: 'curses._CursesWindow', margin: Margin) -> None:
         to_display = min(len(self.lines) - self.file_y, margin.body_lines)
-        for i in range(to_display):
-            line_idx = self.file_y + i
-            line = self.lines[line_idx]
-            x = self.x if line_idx == self.y else 0
-            line = scrolled_line(line, x, curses.COLS)
-            stdscr.insstr(i + margin.header, 0, line)
-        for i in range(to_display, margin.body_lines):
-            stdscr.move(i + margin.header, 0)
-            stdscr.clrtoeol()
 
         for file_hl in self._file_hls:
             file_hl.highlight_until(self.lines, self.file_y + to_display)
 
-        for i in range(self.file_y, self.file_y + to_display):
+        for i in range(to_display):
+            draw_y = i + margin.header
+            l_y = self.file_y + i
+            x = self.x if l_y == self.y else 0
+            line = scrolled_line(self.lines[l_y], x, curses.COLS)
+            stdscr.insstr(draw_y, 0, line)
+
+            l_x = line_x(x, curses.COLS)
+            l_x_max = l_x + curses.COLS
             for file_hl in self._file_hls:
-                for region in file_hl.regions[i]:
-                    self.highlight(
-                        stdscr, margin,
-                        y=i, include_edge=file_hl.include_edge, **region,
-                    )
+                for region in file_hl.regions[l_y]:
+                    if region.x >= l_x_max:
+                        break
+                    elif region.end < l_x:
+                        continue
 
-        if self.select_start is not None:
-            (s_y, s_x), (e_y, e_x) = self._get_selection()
+                    if l_x and region.x <= l_x:
+                        if file_hl.include_edge:
+                            h_s_x = 0
+                        else:
+                            h_s_x = 1
+                    else:
+                        h_s_x = region.x - l_x
 
-            if s_y == e_y:
-                self.highlight(
-                    stdscr, margin,
-                    y=s_y, x=s_x, n=e_x - s_x,
-                    color=HIGHLIGHT, include_edge=True,
-                )
-            else:
-                self.highlight(
-                    stdscr, margin,
-                    y=s_y, x=s_x, n=len(self.lines[s_y]) - s_x + 1,
-                    color=HIGHLIGHT, include_edge=True,
-                )
-                for l_y in range(s_y + 1, e_y):
-                    self.highlight(
-                        stdscr, margin,
-                        y=l_y, x=0, n=len(self.lines[l_y]) + 1,
-                        color=HIGHLIGHT, include_edge=True,
-                    )
-                self.highlight(
-                    stdscr, margin,
-                    y=e_y, x=0, n=e_x,
-                    color=HIGHLIGHT, include_edge=True,
-                )
+                    if region.end >= l_x_max:
+                        if file_hl.include_edge:
+                            h_e_x = curses.COLS
+                        else:
+                            h_e_x = curses.COLS - 1
+                    else:
+                        h_e_x = region.end - l_x
 
-    def highlight(
-            self,
-            stdscr: 'curses._CursesWindow', margin: Margin,
-            *,
-            y: int, x: int, n: int, color: int,
-            include_edge: bool,
-    ) -> None:
-        h_y = y - self.file_y + margin.header
-        if y == self.y:
-            l_x = line_x(self.x, curses.COLS)
-            # TODO: include edge left detection
-            if x < l_x:
-                h_x = 0
-                n -= l_x - x
-            else:
-                h_x = x - l_x
-        else:
-            l_x = 0
-            h_x = x
-        if not include_edge and len(self.lines[y]) > l_x + curses.COLS:
-            h_n = min(curses.COLS - h_x - 1, n)
-        else:
-            h_n = n
-        if (
-                h_y < margin.header or
-                h_y > margin.header + margin.body_lines or
-                h_x >= curses.COLS
-        ):
-            return
-        stdscr.chgat(h_y, h_x, h_n, color)
+                    stdscr.chgat(draw_y, h_s_x, h_e_x - h_s_x, region.attr)
+
+        for i in range(to_display, margin.body_lines):
+            stdscr.move(i + margin.header, 0)
+            stdscr.clrtoeol()

babi/hl/interface.py

@@ -1,21 +1,21 @@
+from typing import NamedTuple
 from typing import Tuple
 
 from babi._types import Protocol
-from babi._types import TypedDict
 from babi.list_spy import SequenceNoSlice
 
 
-class CursesRegion(TypedDict):
+class HL(NamedTuple):
     x: int
-    n: int
-    color: int
+    end: int
+    attr: int
 
 
-CursesRegions = Tuple[CursesRegion, ...]
+HLs = Tuple[HL, ...]
 
 
 class RegionsMapping(Protocol):
-    def __getitem__(self, idx: int) -> CursesRegions: ...
+    def __getitem__(self, idx: int) -> HLs: ...
 
 
 class FileHL(Protocol):

babi/hl/replace.py

@@ -4,8 +4,8 @@ import curses
 from typing import Dict
 from typing import Generator
 
-from babi.hl.interface import CursesRegion
-from babi.hl.interface import CursesRegions
+from babi.hl.interface import HL
+from babi.hl.interface import HLs
 from babi.list_spy import SequenceNoSlice
 
 HIGHLIGHT = curses.A_REVERSE | curses.A_DIM
@@ -15,7 +15,7 @@ class Replace:
     include_edge = True
 
     def __init__(self) -> None:
-        self.regions: Dict[int, CursesRegions] = collections.defaultdict(tuple)
+        self.regions: Dict[int, HLs] = collections.defaultdict(tuple)
 
     def highlight_until(self, lines: SequenceNoSlice, idx: int) -> None:
         """our highlight regions are populated in other ways"""
@@ -24,8 +24,8 @@ class Replace:
         """our highlight regions are populated in other ways"""
 
     @contextlib.contextmanager
-    def region(self, y: int, x: int, n: int) -> Generator[None, None, None]:
-        self.regions[y] = (CursesRegion(x=x, n=n, color=HIGHLIGHT),)
+    def region(self, y: int, x: int, end: int) -> Generator[None, None, None]:
+        self.regions[y] = (HL(x=x, end=end, attr=HIGHLIGHT),)
         try:
             yield
         finally:

babi/hl/selection.py

@@ -0,0 +1,58 @@
+import collections
+import curses
+from typing import Dict
+from typing import Optional
+from typing import Tuple
+
+from babi.hl.interface import HL
+from babi.hl.interface import HLs
+from babi.list_spy import SequenceNoSlice
+
+ATTR = curses.A_REVERSE | curses.A_DIM
+
+
+class Selection:
+    include_edge = True
+
+    def __init__(self) -> None:
+        self.regions: Dict[int, HLs] = collections.defaultdict(tuple)
+        self.start: Optional[Tuple[int, int]] = None
+        self.end: Optional[Tuple[int, int]] = None
+
+    def highlight_until(self, lines: SequenceNoSlice, idx: int) -> None:
+        if self.start is None or self.end is None:
+            return
+
+        (s_y, s_x), (e_y, e_x) = self.get()
+        if s_y == e_y:
+            self.regions[s_y] = (HL(x=s_x, end=e_x, attr=ATTR),)
+        else:
+            self.regions[s_y] = (
+                HL(x=s_x, end=len(lines[s_y]) + 1, attr=ATTR),
+            )
+            for l_y in range(s_y + 1, e_y):
+                self.regions[l_y] = (
+                    HL(x=0, end=len(lines[l_y]) + 1, attr=ATTR),
+                )
+            self.regions[e_y] = (HL(x=0, end=e_x, attr=ATTR),)
+
+    def touch(self, lineno: int) -> None:
+        """our highlight regions are populated in other ways"""
+
+    def get(self) -> Tuple[Tuple[int, int], Tuple[int, int]]:
+        assert self.start is not None and self.end is not None
+        if self.start < self.end:
+            return self.start, self.end
+        else:
+            return self.end, self.start
+
+    def clear(self) -> None:
+        if self.start is not None and self.end is not None:
+            (s_y, _), (e_y, _) = self.get()
+            for l_y in range(s_y, e_y + 1):
+                del self.regions[l_y]
+        self.start = self.end = None
+
+    def set(self, s_y: int, s_x: int, e_y: int, e_x: int) -> None:
+        self.clear()
+        self.start, self.end = (s_y, s_x), (e_y, e_x)

babi/hl/syntax.py

@@ -9,8 +9,8 @@ from babi.highlight import Compiler
 from babi.highlight import Grammars
 from babi.highlight import highlight_line
 from babi.highlight import State
-from babi.hl.interface import CursesRegion
-from babi.hl.interface import CursesRegions
+from babi.hl.interface import HL
+from babi.hl.interface import HLs
 from babi.list_spy import SequenceNoSlice
 from babi.theme import Style
 from babi.theme import Theme
@@ -33,10 +33,10 @@ class FileSyntax:
         self._theme = theme
         self._color_manager = color_manager
 
-        self.regions: List[CursesRegions] = []
+        self.regions: List[HLs] = []
         self._states: List[State] = []
 
-        self._hl_cache: Dict[str, Dict[State, Tuple[State, CursesRegions]]]
+        self._hl_cache: Dict[str, Dict[State, Tuple[State, HLs]]]
         self._hl_cache = {}
 
     def attr(self, style: Style) -> int:
@@ -53,7 +53,7 @@ class FileSyntax:
             state: State,
             line: str,
             i: int,
-    ) -> Tuple[State, CursesRegions]:
+    ) -> Tuple[State, HLs]:
         try:
             return self._hl_cache[line][state]
         except KeyError:
@@ -67,22 +67,21 @@ class FileSyntax:
         new_end = regions[-1]._replace(end=regions[-1].end - 1)
         regions = regions[:-1] + (new_end,)
 
-        regs: List[CursesRegion] = []
+        regs: List[HL] = []
         for r in regions:
             style = self._theme.select(r.scope)
             if style == self._theme.default:
                 continue
 
-            n = r.end - r.start
             attr = self.attr(style)
             if (
                     regs and
-                    regs[-1]['color'] == attr and
-                    regs[-1]['x'] + regs[-1]['n'] == r.start
+                    regs[-1].attr == attr and
+                    regs[-1].end == r.start
             ):
-                regs[-1]['n'] += n
+                regs[-1] = regs[-1]._replace(end=r.end)
             else:
-                regs.append(CursesRegion(x=r.start, n=n, color=attr))
+                regs.append(HL(x=r.start, end=r.end, attr=attr))
 
         dct = self._hl_cache.setdefault(line, {})
         ret = dct[state] = (new_state, tuple(regs))

babi/hl/trailing_whitespace.py

@@ -3,8 +3,8 @@ from typing import List
 from typing import NamedTuple
 
 from babi.color_manager import ColorManager
-from babi.hl.interface import CursesRegion
-from babi.hl.interface import CursesRegions
+from babi.hl.interface import HL
+from babi.hl.interface import HLs
 from babi.list_spy import SequenceNoSlice
 
 
@@ -14,9 +14,9 @@ class FileTrailingWhitespace:
     def __init__(self, color_manager: ColorManager) -> None:
         self._color_manager = color_manager
 
-        self.regions: List[CursesRegions] = []
+        self.regions: List[HLs] = []
 
-    def _trailing_ws(self, line: str) -> CursesRegions:
+    def _trailing_ws(self, line: str) -> HLs:
         if not line:
             return ()
 
@@ -29,7 +29,7 @@ class FileTrailingWhitespace:
         else:
             pair = self._color_manager.raw_color_pair(-1, curses.COLOR_RED)
             attr = curses.color_pair(pair)
-            return (CursesRegion(x=i, n=len(line) - i, color=attr),)
+            return (HL(x=i, end=len(line), attr=attr),)
 
     def highlight_until(self, lines: SequenceNoSlice, idx: int) -> None:
         for i in range(len(self.regions), idx):

babi/horizontal_scrolling.py

@@ -1,10 +1,10 @@
 def line_x(x: int, width: int) -> int:
-    margin = min(width - 3, 6)
     if x + 1 < width:
         return 0
     elif width == 1:
         return x
     else:
+        margin = min(width - 3, 6)
         return (
             width - margin - 2 +
             (x + 1 - width) //
@@ -17,7 +17,7 @@ def scrolled_line(s: str, x: int, width: int) -> str:
     l_x = line_x(x, width)
     if l_x:
         s = f'«{s[l_x + 1:]}'
-        if l_x and len(s) > width:
+        if len(s) > width:
             return f'{s[:width - 1]}»'
         else:
             return s.ljust(width)

babi/screen.py

@@ -293,7 +293,7 @@ class Screen:
         self.status.update(f'{line}, {col} (of {line_count} {lines_word})')
 
     def cut(self) -> None:
-        if self.file.select_start:
+        if self.file.selection.start:
             self.cut_buffer = self.file.cut_selection(self.margin)
             self.cut_selection = True
         else:
@@ -329,6 +329,7 @@ class Screen:
             to_stack.append(action.apply(self.file))
             self.file.scroll_screen_if_needed(self.margin)
             self.status.update(f'{op}: {action.name}')
+            self.file.selection.clear()
 
     def undo(self) -> None:
         self._undo_redo('undo', self.file.undo_stack, self.file.redo_stack)
@@ -360,7 +361,7 @@ class Screen:
             self.save()
             return EditResult.EXIT
         elif response == ':sort':
-            if self.file.select_start:
+            if self.file.selection.start:
                 self.file.sort_selection(self.margin)
             else:
                 self.file.sort(self.margin)

tests/features/syntax_highlight_test.py

@@ -78,3 +78,22 @@ def test_syntax_highlighting_does_not_highlight_arrows(run, tmpdir):
     with run(str(f), term='screen-256color', width=20) as h, and_exit(h):
         h.await_text('loooo')
         h.assert_screen_attr_equals(2, [(243, 40, 0)] * 19 + [(236, 40, 0)])
+
+        h.press('Down')
+        h.press('^E')
+        h.await_text_missing('loooo')
+        expected = [(236, 40, 0)] + [(243, 40, 0)] * 15 + [(236, 40, 0)] * 4
+        h.assert_screen_attr_equals(2, expected)
+
+
+def test_syntax_highlighting_off_screen_does_not_crash(run, tmpdir):
+    f = tmpdir.join('f.demo')
+    f.write(f'"""a"""{"x" * 40}"""b"""')
+
+    with run(str(f), term='screen-256color', width=20) as h, and_exit(h):
+        h.await_text('"""a"""')
+        h.assert_screen_attr_equals(1, [(17, 40, 0)] * 7 + [(236, 40, 0)] * 13)
+        h.press('^E')
+        h.await_text('"""b"""')
+        expected = [(236, 40, 0)] * 11 + [(17, 40, 0)] * 7 + [(236, 40, 0)] * 2
+        h.assert_screen_attr_equals(1, expected)

tests/features/undo_redo_test.py

@@ -127,3 +127,16 @@ def test_undo_redo_causes_scroll(run):
         h.await_cursor_position(x=0, y=1)
         h.press('M-U')
         h.await_cursor_position(x=0, y=4)
+
+
+def test_undo_redo_clears_selection(run, ten_lines):
+    # maintaining the selection across undo/redo is both difficult and not all
+    # that useful.  prior to this it was buggy anyway (a negative selection
+    # indented and then undone would highlight out of bounds)
+    with run(str(ten_lines), width=20) as h, and_exit(h):
+        h.press('S-Down')
+        h.press('Tab')
+        h.await_cursor_position(x=4, y=2)
+        h.press('M-u')
+        h.await_cursor_position(x=0, y=2)
+        h.assert_screen_attr_equals(1, [(-1, -1, 0)] * 20)