Make syntax highlight generation scale better.

The TextArea widget code queries the entires syntax tree for each edit,
using the tree-sitter Query.captures method. This has potential scaling
issues, but such issues are exacerbated by the fact the Query.captures
method scales very poorly with the number of line it is asked to
generate captures for. It appears to be quadratic or something similar I
think - I strongly suspect a bug in tree-sitter or its python bindings.

On my laptop, this makes editing a 25,000 line Python file painfully
unresponsive. A 25,000 lines Python file is large, but not entirely
unreasonable. I actually became aware of this behaviour developing a
LaTeX editor, which showed symptoms after about 1,000 lines.

This change replaces the plain TextArea._highlights dictionary with a
dictonary-like class that lazily performs item access to build
hilghlight information for small blocks of (50) lines at a time. As well
as keeping the TextArea very much more responsive, it will reduce the
average memory requirements for larger documents.

During regression testing, I discovered that the per-line highlights are
not necessarily in the correct (best) order for syntax highlighting. For
example, highlighting within string expressions can lost. So I added
suitable sorting. This required that the snapshots for some tests to be
updated.

Author: paul-ollis

CHANGELOG.md

@@ -61,6 +61,17 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 - Added `Style.has_transparent_foreground` property https://github.com/Textualize/textual/pull/5657
 
 
+### Fixed
+
+- Fixed TextArea's syntax highlighting. Some highlighting details were not being
+  applied. For example, in CSS, the text 'padding: 10px 0;' was shown in a
+  single colour. Now the 'px' appears in a different colour to the rest of the
+  text.
+
+- Fixed a cause of slow editing for syntax highlighed TextArea widgets with
+  large documents.
+
+
 ## [2.1.2] - 2025-02-26
 
 ### Fixed

src/textual/document/_syntax_aware_document.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+from contextlib import contextmanager
+from typing import ContextManager
+
 try:
     from tree_sitter import Language, Node, Parser, Query, Tree
 
@@ -11,6 +14,35 @@
 from textual.document._document import Document, EditResult, Location, _utf8_encode
 
 
+@contextmanager
+def temporary_query_point_range(
+    query: Query,
+    start_point: tuple[int, int] | None,
+    end_point: tuple[int, int] | None,
+) -> ContextManager[None]:
+    """Temporarily change the start and/or end point for a tree-sitter Query.
+
+    Args:
+        query: The tree-sitter Query.
+        start_point: The (row, column byte) to start the query at.
+        end_point: The (row, column byte) to end the query at.
+    """
+    # Note: Although not documented for the tree-sitter Python API, an
+    # end-point of (0, 0) means 'end of document'.
+    default_point_range = [(0, 0), (0, 0)]
+
+    point_range = list(default_point_range)
+    if start_point is not None:
+        point_range[0] = start_point
+    if end_point is not None:
+        point_range[1] = end_point
+    query.set_point_range(point_range)
+    try:
+        yield None
+    finally:
+        query.set_point_range(default_point_range)
+
+
 class SyntaxAwareDocumentError(Exception):
     """General error raised when SyntaxAwareDocument is used incorrectly."""
 
@@ -82,14 +114,8 @@ def query_syntax_tree(
         Returns:
             A tuple containing the nodes and text captured by the query.
         """
-        captures_kwargs = {}
-        if start_point is not None:
-            captures_kwargs["start_point"] = start_point
-        if end_point is not None:
-            captures_kwargs["end_point"] = end_point
-
-        captures = query.captures(self._syntax_tree.root_node, **captures_kwargs)
-        return captures
+        with temporary_query_point_range(query, start_point, end_point):
+            return query.captures(self._syntax_tree.root_node)
 
     def replace_range(self, start: Location, end: Location, text: str) -> EditResult:
         """Replace text at the given range.

src/textual/widgets/_text_area.py

@@ -89,6 +89,105 @@ class LanguageDoesNotExist(Exception):
     """
 
 
+class HighlightMap:
+    """Lazy evaluated pseudo dictionary mapping lines to highlight information.
+
+    This allows TextArea syntax highlighting to scale.
+
+    Args:
+        text_area_widget: The associated `TextArea` widget.
+    """
+
+    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] = {}
+
+        # A mapping from line index to a list of Highlight instances.
+        self._highlights: LineToHighlightsMap = defaultdict(list)
+        self.reset()
+
+    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})
+
+    @property
+    def document(self) -> DocumentBase:
+        """The text document being highlighted."""
+        return self.text_area_widget.document
+
+    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, line_range: range) -> None:
+        """Build part of the highlight 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)
+        captures = self.document.query_syntax_tree(
+            self.text_area_widget._highlight_query,
+            start_point=start_point,
+            end_point=end_point,
+        )
+        for highlight_name, nodes in captures.items():
+            for node in nodes:
+                node_start_row, node_start_column = node.start_point
+                node_end_row, node_end_column = node.end_point
+                if node_start_row == node_end_row:
+                    highlight = node_start_column, node_end_column, highlight_name
+                    highlights[node_start_row].append(highlight)
+                else:
+                    # Add the first line of the node range
+                    highlights[node_start_row].append(
+                        (node_start_column, None, highlight_name)
+                    )
+
+                    # Add the middle lines - entire row of this node is highlighted
+                    for node_row in range(node_start_row + 1, node_end_row):
+                        highlights[node_row].append((0, None, highlight_name))
+
+                    # Add the last line of the node range
+                    highlights[node_end_row].append(
+                        (0, node_end_column, highlight_name)
+                    )
+
+        # The highlights for each line need to be sorted. Each highlight is of
+        # the form:
+        #
+        #     a, b, highlight-name
+        #
+        # Where a is a number and b is a number or ``None``. These highlights need
+        # 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
+            if b is None:
+                max_range_ind = 0
+                b = a
+            return a, max_range_ind, a - b
+
+        for line_index in line_range:
+            line_highlights = highlights.get(line_index, []).sort(key=sort_key)
+
+
 @dataclass
 class TextAreaLanguage:
     """A container for a language which has been registered with the TextArea."""
@@ -472,15 +571,15 @@ def __init__(
         cursor is currently at. If the cursor is at a bracket, or there's no matching
         bracket, this will be `None`."""
 
-        self._highlights: dict[int, list[Highlight]] = defaultdict(list)
-        """Mapping line numbers to the set of highlights for that line."""
-
         self._highlight_query: "Query | None" = None
         """The query that's currently being used for highlighting."""
 
         self.document: DocumentBase = Document(text)
         """The document this widget is currently editing."""
 
+        self._highlights: HighlightMap = HighlightMap(self)
+        """Mapping line numbers to the set of highlights for that line."""
+
         self.wrapped_document: WrappedDocument = WrappedDocument(self.document)
         """The wrapped view of the document."""
 
@@ -609,35 +708,10 @@ def check_consume_key(self, key: str, character: str | None = None) -> bool:
         return character is not None and character.isprintable()
 
     def _build_highlight_map(self) -> None:
-        """Query the tree for ranges to highlights, and update the internal highlights mapping."""
-        highlights = self._highlights
-        highlights.clear()
-        if not self._highlight_query:
-            return
-
-        captures = self.document.query_syntax_tree(self._highlight_query)
-        for highlight_name, nodes in captures.items():
-            for node in nodes:
-                node_start_row, node_start_column = node.start_point
-                node_end_row, node_end_column = node.end_point
-
-                if node_start_row == node_end_row:
-                    highlight = (node_start_column, node_end_column, highlight_name)
-                    highlights[node_start_row].append(highlight)
-                else:
-                    # Add the first line of the node range
-                    highlights[node_start_row].append(
-                        (node_start_column, None, highlight_name)
-                    )
+        """Reset the lazily evaluated highlight map."""
 
-                    # Add the middle lines - entire row of this node is highlighted
-                    for node_row in range(node_start_row + 1, node_end_row):
-                        highlights[node_row].append((0, None, highlight_name))
-
-                    # Add the last line of the node range
-                    highlights[node_end_row].append(
-                        (0, node_end_column, highlight_name)
-                    )
+        if self._highlight_query:
+            self._highlights.reset()
 
     def _watch_has_focus(self, focus: bool) -> None:
         self._cursor_visible = focus