Restored callback approach but retained LRU caching of pt styles.

Author: kmvanbrunt

CHANGELOG.md

@@ -153,6 +153,8 @@ 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,13 +526,10 @@ def __init__(
         self._persistent_history_length = persistent_history_length
         self._initialize_history(persistent_history_file)
 
-        # 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, ...] = ()
+        # Cache for prompt_toolkit completion menu styles
+        self.pt_style: PtStyle
+        self.update_pt_style()
+        ru.register_theme_update_callback(self.update_pt_style)
 
         # Create the main PromptSession
         self.bottom_toolbar = bottom_toolbar
@@ -727,39 +724,35 @@ def _should_continue_multiline(self) -> bool:
                 # No macro found or already processed. The statement is complete.
                 return False
 
-    def _get_pt_style(self) -> PtStyle:
-        """Return the prompt_toolkit style synchronized with the application theme."""
+    def update_pt_style(self) -> None:
+        """Update the cached prompt_toolkit style."""
         theme = ru.get_theme()
-
-        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,
+        rich_menu_style = theme.styles.get(Cmd2Style.COMPLETION_MENU, Style.null())
+        rich_completion_style = theme.styles.get(Cmd2Style.COMPLETION_MENU_COMPLETION, Style.null())
+        rich_current_style = theme.styles.get(Cmd2Style.COMPLETION_MENU_CURRENT, Style.null())
+        rich_meta_style = theme.styles.get(Cmd2Style.COMPLETION_MENU_META, Style.null())
+        rich_meta_current_style = theme.styles.get(Cmd2Style.COMPLETION_MENU_META_CURRENT, Style.null())
+
+        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,
+            }
         )
 
-        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 _get_pt_style(self) -> "PtStyle":
+        """Return the cached prompt_toolkit style."""
+        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,6 +1,7 @@
 """Utilities for integrating prompt_toolkit with cmd2."""
 
 import re
+import weakref
 from collections.abc import (
     Callable,
     Iterable,
@@ -263,6 +264,18 @@ 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."""
 
@@ -277,29 +290,18 @@ def __init__(
         super().__init__()
         self.cmd_app = cmd_app
 
-        # Cache key used to detect when theme styles have changed
-        self._style_key: tuple[Style, ...] = ()
+        _lexers.add(self)
         self.set_colors()
 
     def set_colors(self) -> None:
-        """Synchronize lexer colors with the application theme."""
+        """Update colors from the current rich theme."""
+        # Retrieve styles dynamically from the current theme
         theme = ru.get_theme()
-
-        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)
+        self.command_color = rich_to_pt_style(theme.styles.get(Cmd2Style.LEXER_COMMAND, Style.null()))
+        self.alias_color = rich_to_pt_style(theme.styles.get(Cmd2Style.LEXER_ALIAS, Style.null()))
+        self.macro_color = rich_to_pt_style(theme.styles.get(Cmd2Style.LEXER_MACRO, Style.null()))
+        self.flag_color = rich_to_pt_style(theme.styles.get(Cmd2Style.LEXER_FLAG, Style.null()))
+        self.argument_color = rich_to_pt_style(theme.styles.get(Cmd2Style.LEXER_ARGUMENT, Style.null()))
 
     def lex_document(self, document: Document) -> Callable[[int], Any]:
         """Lex the document."""
@@ -308,8 +310,6 @@ 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,6 +4,7 @@
 import re
 import sys
 from collections.abc import (
+    Callable,
     Iterator,
     Mapping,
 )
@@ -309,6 +310,15 @@ 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."""
@@ -351,6 +361,10 @@ 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,8 +279,7 @@ 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 the lexer uses current theme values."""
-        mock_cmd_app.all_commands = ["help"]
+        """Test that changing the theme updates active lexers."""
         lexer = pt_utils.Cmd2Lexer(cast(Any, mock_cmd_app))
 
         # Get the old color for command
@@ -296,13 +295,7 @@ def test_lexer_set_theme_runtime_update(self, mock_cmd_app):
         try:
             ru.set_theme(new_styles)
 
-            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")]
+            # Now verify the lexer's color was updated
             assert lexer.command_color != old_color
             assert "ansired" in lexer.command_color
             assert "ansiblack" in lexer.command_color
@@ -810,3 +803,12 @@ 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,3 +349,26 @@ 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()