Update with suggestions by Darren Burns.

Author: paul-ollis

src/textual/widgets/_text_area.py

@@ -101,48 +101,47 @@ class HighlightMap:
 
     BLOCK_SIZE = 50
 
-    def __init__(self, text_area_widget: widgets.TextArea):
-        self.text_area_widget: widgets.TextArea = text_area_widget
-        self.uncovered_lines: dict[int, range] = {}
+    def __init__(self, text_area: TextArea):
+        self.text_area: TextArea = text_area
+        """The text area associated with this highlight map."""
 
-        # A mapping from line index to a list of Highlight instances.
-        self._highlights: LineToHighlightsMap = defaultdict(list)
-        self.reset()
+        self._highlighted_blocks: set[int] = set()
+        """The set of blocks that have been highlighted. Each block covers BLOCK_SIZE
+        lines.
+        """
+
+        self._highlights: dict[int, list[Highlight]] = defaultdict(list)
+        """A mapping from line index to a list of Highlight instances."""
 
     def reset(self) -> None:
         """Reset so that future lookups rebuild the highlight map."""
         self._highlights.clear()
-        line_count = self.document.line_count
-        uncovered_lines = self.uncovered_lines
-        uncovered_lines.clear()
-        i = end_range = 0
-        for i in range(0, line_count, self.BLOCK_SIZE):
-            end_range = min(i + self.BLOCK_SIZE, line_count)
-            line_range = range(i, end_range)
-            uncovered_lines.update({j: line_range for j in line_range})
-        if end_range < line_count:
-            line_range = range(i, line_count)
-            uncovered_lines.update({j: line_range for j in line_range})
+        self._highlighted_blocks.clear()
 
     @property
     def document(self) -> DocumentBase:
         """The text document being highlighted."""
-        return self.text_area_widget.document
+        return self.text_area.document
+
+    def __getitem__(self, index: int) -> list[Highlight]:
+        block_index = index // self.BLOCK_SIZE
+        if block_index not in self._highlighted_blocks:
+            self._highlighted_blocks.add(block_index)
+            self._build_part_of_highlight_map(block_index * self.BLOCK_SIZE)
+        return self._highlights[index]
 
-    def __getitem__(self, idx: int) -> list[text_area.Highlight]:
-        if idx in self.uncovered_lines:
-            self._build_part_of_highlight_map(self.uncovered_lines[idx])
-        return self._highlights[idx]
+    def _build_part_of_highlight_map(self, start_index: int) -> None:
+        """Build part of the highlight map.
 
-    def _build_part_of_highlight_map(self, line_range: range) -> None:
-        """Build part of the highlight map."""
+        Args:
+            start_index: The start of the block of line for which to build the map.
+        """
         highlights = self._highlights
-        for line_index in line_range:
-            self.uncovered_lines.pop(line_index)
-        start_point = (line_range[0], 0)
-        end_point = (line_range[-1] + 1, 0)
+        start_point = (start_index, 0)
+        end_index = min(self.document.line_count, start_index + self.BLOCK_SIZE)
+        end_point = (end_index, 0)
         captures = self.document.query_syntax_tree(
-            self.text_area_widget._highlight_query,
+            self.text_area._highlight_query,
             start_point=start_point,
             end_point=end_point,
         )
@@ -160,8 +159,9 @@ def _build_part_of_highlight_map(self, line_range: range) -> None:
                     )
 
                     # Add the middle lines - entire row of this node is highlighted
+                    middle_highlight = (0, None, highlight_name)
                     for node_row in range(node_start_row + 1, node_end_row):
-                        highlights[node_row].append((0, None, highlight_name))
+                        highlights[node_row].append(middle_highlight)
 
                     # Add the last line of the node range
                     highlights[node_end_row].append(
@@ -177,16 +177,16 @@ def _build_part_of_highlight_map(self, line_range: range) -> None:
         # to be sorted in ascending order of ``a``. When two highlights have the same
         # value of ``a`` then the one with the larger a--b range comes first, with ``None``
         # being considered larger than any number.
-        def sort_key(hl) -> tuple[int, int, int]:
-            a, b, _ = hl
-            max_range_ind = 1
+        def sort_key(highlight: Highlight) -> tuple[int, int, int]:
+            a, b, _ = highlight
+            max_range_index = 1
             if b is None:
-                max_range_ind = 0
+                max_range_index = 0
                 b = a
-            return a, max_range_ind, a - b
+            return a, max_range_index, a - b
 
-        for line_index in line_range:
-            line_highlights = highlights.get(line_index, []).sort(key=sort_key)
+        for line_index in range(start_index, end_index):
+            highlights.get(line_index, []).sort(key=sort_key)
 
 
 @dataclass
@@ -708,7 +708,7 @@ def check_consume_key(self, key: str, character: str | None = None) -> bool:
         # Otherwise we capture all printable keys
         return character is not None and character.isprintable()
 
-    def _build_highlight_map(self) -> None:
+    def _reset_highlights(self) -> None:
         """Reset the lazily evaluated highlight map."""
 
         if self._highlight_query:
@@ -1034,7 +1034,7 @@ def _set_document(self, text: str, language: str | None) -> None:
         self.document = document
         self.wrapped_document = WrappedDocument(document, tab_width=self.indent_width)
         self.navigator = DocumentNavigator(self.wrapped_document)
-        self._build_highlight_map()
+        self._reset_highlights()
         self.move_cursor((0, 0))
         self._rewrap_and_refresh_virtual_size()
 
@@ -1447,7 +1447,7 @@ def edit(self, edit: Edit) -> EditResult:
 
         self._refresh_size()
         edit.after(self)
-        self._build_highlight_map()
+        self._reset_highlights()
         self.post_message(self.Changed(self))
         return result
 
@@ -1510,7 +1510,7 @@ def _undo_batch(self, edits: Sequence[Edit]) -> None:
         self._refresh_size()
         for edit in reversed(edits):
             edit.after(self)
-        self._build_highlight_map()
+        self._reset_highlights()
         self.post_message(self.Changed(self))
 
     def _redo_batch(self, edits: Sequence[Edit]) -> None:
@@ -1558,7 +1558,7 @@ def _redo_batch(self, edits: Sequence[Edit]) -> None:
         self._refresh_size()
         for edit in edits:
             edit.after(self)
-        self._build_highlight_map()
+        self._reset_highlights()
         self.post_message(self.Changed(self))
 
     async def _on_key(self, event: events.Key) -> None: