Optimize query_one

Author: willmcgugan

src/textual/css/parse.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import dataclasses
+import re
 from functools import lru_cache
 from typing import Iterable, Iterator, NoReturn
 
@@ -16,7 +17,13 @@
     SelectorType,
 )
 from textual.css.styles import Styles
-from textual.css.tokenize import Token, tokenize, tokenize_declarations, tokenize_values
+from textual.css.tokenize import (
+    IDENTIFIER,
+    Token,
+    tokenize,
+    tokenize_declarations,
+    tokenize_values,
+)
 from textual.css.tokenizer import ReferencedBy, UnexpectedEnd
 from textual.css.types import CSSLocation, Specificity3
 from textual.suggestions import get_suggestion
@@ -33,6 +40,20 @@
     "nested": (SelectorType.NESTED, (0, 0, 0)),
 }
 
+RE_IDENTIFIER = re.compile(r"\#" + IDENTIFIER)
+
+
+def is_id_selector(selector: str) -> bool:
+    """Is the selector an ID selector, i.e. "#foo"?
+
+    Args:
+        selector: A CSS selector.
+
+    Returns:
+        `True` if the selector is a simple ID selector, otherwise `False`.
+    """
+    return RE_IDENTIFIER.fullmatch(selector) is not None
+
 
 def _add_specificity(
     specificity1: Specificity3, specificity2: Specificity3

src/textual/dom.py

@@ -40,7 +40,7 @@
 from textual.css.constants import VALID_DISPLAY, VALID_VISIBILITY
 from textual.css.errors import DeclarationError, StyleValueError
 from textual.css.match import match
-from textual.css.parse import parse_declarations, parse_selectors
+from textual.css.parse import is_id_selector, parse_declarations, parse_selectors
 from textual.css.query import InvalidQueryFormat, NoMatches, TooManyMatches, WrongType
 from textual.css.styles import RenderStyles, Styles
 from textual.css.tokenize import IDENTIFIER
@@ -49,7 +49,7 @@
 from textual.reactive import Reactive, ReactiveError, _Mutated, _watch
 from textual.style import Style as VisualStyle
 from textual.timer import Timer
-from textual.walk import walk_breadth_first, walk_depth_first
+from textual.walk import walk_breadth_first, walk_breadth_search_id, walk_depth_first
 from textual.worker_manager import WorkerManager
 
 if TYPE_CHECKING:
@@ -1466,6 +1466,26 @@ def query_one(
         else:
             query_selector = selector.__name__
 
+        if is_id_selector(query_selector):
+            cache_key = (base_node._nodes._updates, query_selector, expect_type)
+            cached_result = base_node._query_one_cache.get(cache_key)
+            if cached_result is not None:
+                return cached_result
+            if (
+                node := walk_breadth_search_id(
+                    base_node, query_selector[1:], with_root=False
+                )
+            ) is not None:
+                if expect_type is not None and not isinstance(node, expect_type):
+                    raise WrongType(
+                        f"Node matching {query_selector!r} is the wrong type; expected type {expect_type.__name__!r}, found {node}"
+                    )
+                base_node._query_one_cache[cache_key] = node
+                return node
+            raise NoMatches(
+                f"No nodes match {query_selector!r} on {base_node!r} {list(base_node._nodes)}"
+            )
+
         try:
             selector_set = parse_selectors(query_selector)
         except TokenError:
@@ -1492,7 +1512,7 @@ def query_one(
                 base_node._query_one_cache[cache_key] = node
             return node
 
-        raise NoMatches(f"No nodes match {selector!r} on {base_node!r}")
+        raise NoMatches(f"No nodes match {query_selector!r} on {base_node!r}")
 
     if TYPE_CHECKING:
 
@@ -1572,7 +1592,7 @@ def query_exactly_one(
                 base_node._query_one_cache[cache_key] = node
             return node
 
-        raise NoMatches(f"No nodes match {selector!r} on {base_node!r}")
+        raise NoMatches(f"No nodes match {query_selector!r} on {base_node!r}")
 
     if TYPE_CHECKING:
 

src/textual/walk.py

@@ -137,3 +137,33 @@ def walk_breadth_first(
         if isinstance(node, check_type):
             yield node
         extend(node._nodes)
+
+
+def walk_breadth_search_id(
+    root: DOMNode, node_id: str, *, with_root: bool = True
+) -> DOMNode | None:
+    """Special case to walk breadth first searching for a node with a given id.
+
+    This is more efficient that [walk_breadth_first][textual.walk.walk_breadth_first] for this special case, as it can use an index.
+
+    Args:
+        root: The root node (starting point).
+        node_id: Node id to search for.
+        with_root: Consider the root node? If the root has the node id, then return it.
+
+    Returns:
+        A DOMNode if a node was found, otherwise `None`.
+    """
+
+    if with_root and root.id == node_id:
+        return root
+
+    queue: deque[DOMNode] = deque()
+    queue.append(root)
+
+    while queue:
+        node = queue.popleft()
+        if (found_node := node._nodes._get_by_id(node_id)) is not None:
+            return found_node
+        queue.extend(node._nodes)
+    return None