Sync prompt-toolkit styles on demand using caching and LRU-backed lookups instead of callbacks

Author: kmvanbrunt

CHANGELOG.md

@@ -153,8 +153,6 @@ prompt is displayed.
     - Added `Cmd2ArgumentParser.output_to()` context manager to temporarily set the output stream
       during `argparse` operations. This is helpful for directing output for functions like
       `parse_args()`, which default to `sys.stdout` and lack a `file` argument.
-    - Added `cmd2.rich_utils.register_theme_update_callback` function to register callback functions
-      to get called whenever `cmd2.rich_utils.set_theme` is called
     - Added ability to customize `prompt-toolkit` completion menu colors by overriding the following
       fields in the `cmd2` theme:
         - `Cmd2Style.COMPLETION_MENU` - Base style for the entire completion menu container (sets

cmd2/cmd2.py

@@ -526,10 +526,13 @@ def __init__(
         self._persistent_history_length = persistent_history_length
         self._initialize_history(persistent_history_file)
 
-        # Cache for prompt_toolkit completion menu styles
-        self.pt_style: PtStyle
-        self.update_pt_style()
-        ru.register_theme_update_callback(self.update_pt_style)
+        # Styles used in prompt_toolkit elements like completion menus.
+        # This is initialized when first needed for rendering and is only
+        # updated when the application theme changes. self._pt_style_key
+        # is a tuple of the Rich styles used to build self._pt_style and
+        # acts as a cache key to detect when the theme has changed.
+        self._pt_style: PtStyle | None = None
+        self._pt_style_key: tuple[Style, ...] = ()
 
         # Create the main PromptSession
         self.bottom_toolbar = bottom_toolbar
@@ -724,35 +727,39 @@ def _should_continue_multiline(self) -> bool:
                 # No macro found or already processed. The statement is complete.
                 return False
 
-    def update_pt_style(self) -> None:
-        """Update the cached prompt_toolkit style."""
+    def _get_pt_style(self) -> PtStyle:
+        """Return the prompt_toolkit style synchronized with the application theme."""
         theme = ru.get_theme()
-        rich_menu_style = theme.styles.get(Cmd2Style.COMPLETION_MENU, "")
-        rich_completion_style = theme.styles.get(Cmd2Style.COMPLETION_MENU_COMPLETION, "")
-        rich_current_style = theme.styles.get(Cmd2Style.COMPLETION_MENU_CURRENT, "")
-        rich_meta_style = theme.styles.get(Cmd2Style.COMPLETION_MENU_META, "")
-        rich_meta_current_style = theme.styles.get(Cmd2Style.COMPLETION_MENU_META_CURRENT, "")
-
-        menu_style = rich_to_pt_style(rich_menu_style)
-        completion_style = rich_to_pt_style(rich_completion_style)
-        current_style = rich_to_pt_style(rich_current_style)
-        meta_style = rich_to_pt_style(rich_meta_style)
-        meta_current_style = rich_to_pt_style(rich_meta_current_style)
-
-        self.pt_style = PtStyle.from_dict(
-            {
-                "completion-menu": menu_style,
-                "completion-menu.completion": completion_style,
-                "completion-menu.completion.current": current_style,
-                "completion-menu.meta.completion": meta_style,
-                "completion-menu.meta.completion.current": meta_current_style,
-                "completion-menu.multi-column-meta": meta_current_style,
-            }
+
+        completion_menu = theme.styles.get(Cmd2Style.COMPLETION_MENU, Style.null())
+        completion_menu_completion = theme.styles.get(Cmd2Style.COMPLETION_MENU_COMPLETION, Style.null())
+        completion_menu_current = theme.styles.get(Cmd2Style.COMPLETION_MENU_CURRENT, Style.null())
+        completion_menu_meta = theme.styles.get(Cmd2Style.COMPLETION_MENU_META, Style.null())
+        completion_menu_meta_current = theme.styles.get(Cmd2Style.COMPLETION_MENU_META_CURRENT, Style.null())
+
+        current_key = (
+            completion_menu,
+            completion_menu_completion,
+            completion_menu_current,
+            completion_menu_meta,
+            completion_menu_meta_current,
         )
 
-    def _get_pt_style(self) -> "PtStyle":
-        """Return the cached prompt_toolkit style."""
-        return self.pt_style
+        if self._pt_style is None or current_key != self._pt_style_key:
+            self._pt_style_key = current_key
+
+            self._pt_style = PtStyle.from_dict(
+                {
+                    "completion-menu": rich_to_pt_style(completion_menu),
+                    "completion-menu.completion": rich_to_pt_style(completion_menu_completion),
+                    "completion-menu.completion.current": rich_to_pt_style(completion_menu_current),
+                    "completion-menu.meta.completion": rich_to_pt_style(completion_menu_meta),
+                    "completion-menu.meta.completion.current": rich_to_pt_style(completion_menu_meta_current),
+                    "completion-menu.multi-column-meta": rich_to_pt_style(completion_menu_meta_current),
+                }
+            )
+
+        return self._pt_style
 
     def _create_main_session(self, auto_suggest: bool, completekey: str) -> PromptSession[str]:
         """Create and return the main PromptSession for the application.

cmd2/pt_utils.py

@@ -1,11 +1,11 @@
 """Utilities for integrating prompt_toolkit with cmd2."""
 
 import re
-import weakref
 from collections.abc import (
     Callable,
     Iterable,
 )
+from functools import lru_cache
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -75,6 +75,7 @@ def pt_filter_style(text: str | ANSI) -> str | ANSI:
     return text if isinstance(text, ANSI) else ANSI(text)
 
 
+@lru_cache(maxsize=256)
 def rich_to_pt_color(color: "Color | None") -> str:
     """Convert a rich Color object to a prompt_toolkit color string."""
     if not color or color.is_default:
@@ -90,6 +91,7 @@ def rich_to_pt_color(color: "Color | None") -> str:
     return f"#{c.red:02x}{c.green:02x}{c.blue:02x}"
 
 
+@lru_cache(maxsize=1024)
 def rich_to_pt_style(rich_style: StyleType) -> str:
     """Convert a rich Style object to a prompt_toolkit style string."""
     if not rich_style:
@@ -115,10 +117,8 @@ def rich_to_pt_style(rich_style: StyleType) -> str:
     if rich_style.blink is not None:
         parts.append("blink" if rich_style.blink else "noblink")
     if rich_style.reverse is not None:
-        # prompt-toolkit uses 'reverse'
         parts.append("reverse" if rich_style.reverse else "noreverse")
     if rich_style.conceal is not None:
-        # prompt-toolkit uses 'hidden' for Rich's 'conceal'
         parts.append("hidden" if rich_style.conceal else "nohidden")
     return " ".join(parts)
 
@@ -263,18 +263,6 @@ def clear(self) -> None:
         self._loaded_strings.clear()
 
 
-_lexers: "weakref.WeakSet[Cmd2Lexer]" = weakref.WeakSet()
-
-
-def _update_lexer_colors() -> None:
-    """Update colors for all active lexers."""
-    for lexer in _lexers:
-        lexer.set_colors()
-
-
-ru.register_theme_update_callback(_update_lexer_colors)
-
-
 class Cmd2Lexer(Lexer):
     """Lexer that highlights cmd2 command names, aliases, and macros."""
 
@@ -289,18 +277,29 @@ def __init__(
         super().__init__()
         self.cmd_app = cmd_app
 
-        _lexers.add(self)
+        # Cache key used to detect when theme styles have changed
+        self._style_key: tuple[Style, ...] = ()
         self.set_colors()
 
     def set_colors(self) -> None:
-        """Update colors from the current rich theme."""
-        # Retrieve styles dynamically from the current theme
+        """Synchronize lexer colors with the application theme."""
         theme = ru.get_theme()
-        self.command_color = rich_to_pt_style(theme.styles.get(Cmd2Style.LEXER_COMMAND, ""))
-        self.alias_color = rich_to_pt_style(theme.styles.get(Cmd2Style.LEXER_ALIAS, ""))
-        self.macro_color = rich_to_pt_style(theme.styles.get(Cmd2Style.LEXER_MACRO, ""))
-        self.flag_color = rich_to_pt_style(theme.styles.get(Cmd2Style.LEXER_FLAG, ""))
-        self.argument_color = rich_to_pt_style(theme.styles.get(Cmd2Style.LEXER_ARGUMENT, ""))
+
+        command_style = theme.styles.get(Cmd2Style.LEXER_COMMAND, Style.null())
+        alias_style = theme.styles.get(Cmd2Style.LEXER_ALIAS, Style.null())
+        macro_style = theme.styles.get(Cmd2Style.LEXER_MACRO, Style.null())
+        flag_style = theme.styles.get(Cmd2Style.LEXER_FLAG, Style.null())
+        argument_style = theme.styles.get(Cmd2Style.LEXER_ARGUMENT, Style.null())
+
+        current_key = (command_style, alias_style, macro_style, flag_style, argument_style)
+
+        if current_key != self._style_key:
+            self._style_key = current_key
+            self.command_color = rich_to_pt_style(command_style)
+            self.alias_color = rich_to_pt_style(alias_style)
+            self.macro_color = rich_to_pt_style(macro_style)
+            self.flag_color = rich_to_pt_style(flag_style)
+            self.argument_color = rich_to_pt_style(argument_style)
 
     def lex_document(self, document: Document) -> Callable[[int], Any]:
         """Lex the document."""
@@ -309,6 +308,8 @@ def lex_document(self, document: Document) -> Callable[[int], Any]:
         exclude_tokens.update(self.cmd_app.statement_parser.terminators)
         arg_pattern = re.compile(r'(\s+)|(--?[^\s\'"]+)|("[^"]*"?|\'[^\']*\'?)|([^\s\'"]+)')
 
+        self.set_colors()
+
         def highlight_args(text: str, tokens: list[tuple[str, str]]) -> None:
             """Highlight arguments in a string."""
             for m in arg_pattern.finditer(text):

cmd2/rich_utils.py

@@ -4,7 +4,6 @@
 import re
 import sys
 from collections.abc import (
-    Callable,
     Iterator,
     Mapping,
 )
@@ -310,15 +309,6 @@ def __cmd2_argparse_help__(self, formatter: Cmd2HelpFormatter) -> Group:
 # The application-wide theme. Use get_theme() and set_theme() to access it.
 _APP_THEME: Theme | None = None
 
-# Callbacks to be executed when the theme is updated
-_theme_update_callbacks: list[Callable[[], None]] = []
-
-
-def register_theme_update_callback(callback: Callable[[], None]) -> None:
-    """Register a callback to be executed when the theme is updated."""
-    if callback not in _theme_update_callbacks:
-        _theme_update_callbacks.append(callback)
-
 
 def get_theme() -> Theme:
     """Get the application-wide theme. Initializes it on the first call."""
@@ -361,10 +351,6 @@ def set_theme(styles: Mapping[str, StyleType] | None = None) -> None:
     for name in Cmd2HelpFormatter.styles.keys() & theme.styles.keys():
         Cmd2HelpFormatter.styles[name] = theme.styles[name]
 
-    # Notify callbacks that the theme has been updated
-    for callback in _theme_update_callbacks:
-        callback()
-
 
 def _create_default_theme() -> Theme:
     """Create a default theme for the application.

tests/test_pt_utils.py

@@ -279,7 +279,8 @@ def test_lex_document_multiline(self, mock_cmd_app):
         assert tokens1 == [("fg:ansiyellow bg:default", "help")]
 
     def test_lexer_set_theme_runtime_update(self, mock_cmd_app):
-        """Test that changing the theme updates active lexers."""
+        """Test that the lexer uses current theme values."""
+        mock_cmd_app.all_commands = ["help"]
         lexer = pt_utils.Cmd2Lexer(cast(Any, mock_cmd_app))
 
         # Get the old color for command
@@ -295,7 +296,13 @@ def test_lexer_set_theme_runtime_update(self, mock_cmd_app):
         try:
             ru.set_theme(new_styles)
 
-            # Now verify the lexer's color was updated
+            line = "help"
+            document = Document(line)
+            get_line = lexer.lex_document(document)
+            tokens = get_line(0)
+
+            # Now verify the updated colors were used
+            assert tokens == [("fg:ansired bg:ansiblack", "help")]
             assert lexer.command_color != old_color
             assert "ansired" in lexer.command_color
             assert "ansiblack" in lexer.command_color
@@ -803,12 +810,3 @@ def test_rich_to_pt_style_nohidden_conceal(self):
         style = Style(conceal=False)
         pt_style = pt_utils.rich_to_pt_style(style)
         assert "nohidden" in pt_style
-
-
-def test_update_lexer_colors() -> None:
-    mock_lexer = Mock()
-    pt_utils._lexers.add(mock_lexer)
-
-    pt_utils._update_lexer_colors()
-
-    mock_lexer.set_colors.assert_called_once()

tests/test_rich_utils.py

@@ -349,26 +349,3 @@ def side_effect(color: bool, **kwargs: Any) -> None:
     assert mock_set_color.call_count == 2
     mock_set_color.assert_any_call(True, file=sys.stdout)
     mock_set_color.assert_any_call(True)
-
-
-def test_register_theme_update_callback() -> None:
-    # Clear callbacks for a clean state
-    ru._theme_update_callbacks.clear()
-
-    # Define a dummy callback
-    def my_callback() -> None:
-        pass
-
-    ru.register_theme_update_callback(my_callback)
-    assert my_callback in ru._theme_update_callbacks
-
-    # Test that registering the same callback again doesn't duplicate it
-    ru.register_theme_update_callback(my_callback)
-    assert len(ru._theme_update_callbacks) == 1
-
-    # Test that set_theme calls the callback
-    mock_callback = mock.Mock()
-    ru.register_theme_update_callback(mock_callback)
-
-    ru.set_theme()
-    mock_callback.assert_called_once()