Update Cmd2Lexer to recognize shortcuts and color them like a command

Author: tleonhardt

.pre-commit-config.yaml

@@ -11,7 +11,7 @@ repos:
         exclude: ^examples/transcripts/
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: "v0.15.0"
+    rev: "v0.15.1"
     hooks:
       - id: ruff-format
         args: [--config=ruff.toml]

cmd2/pt_utils.py

@@ -37,7 +37,11 @@
 class Cmd2Completer(Completer):
     """Completer that delegates to cmd2's completion logic."""
 
-    def __init__(self, cmd_app: 'Cmd', custom_settings: utils.CustomCompletionSettings | None = None) -> None:
+    def __init__(
+        self,
+        cmd_app: 'Cmd',
+        custom_settings: utils.CustomCompletionSettings | None = None,
+    ) -> None:
         """Initialize prompt_toolkit based completer class."""
         self.cmd_app = cmd_app
         self.custom_settings = custom_settings
@@ -155,10 +159,31 @@ def store_string(self, string: str) -> None:
 class Cmd2Lexer(Lexer):
     """Lexer that highlights cmd2 command names, aliases, and macros."""
 
-    def __init__(self, cmd_app: 'Cmd') -> None:
-        """Initialize the lexer."""
+    def __init__(
+        self,
+        cmd_app: 'Cmd',
+        command_color: str = 'ansigreen',
+        alias_color: str = 'ansicyan',
+        macro_color: str = 'ansimagenta',
+        flag_color: str = 'ansired',
+        argument_color: str = 'ansiyellow',
+    ) -> None:
+        """Initialize the Lexer.
+
+        :param cmd_app: cmd2.Cmd instance
+        :param command_color: color to use for commands, defaults to 'ansigreen'
+        :param alias_color: color to use for aliases, defaults to 'ansicyan'
+        :param macro_color: color to use for macros, defaults to 'ansimagenta'
+        :param flag_color: color to use for flags, defaults to 'ansired'
+        :param argument_color: color to use for arguments, defaults to 'ansiyellow'
+        """
         super().__init__()
         self.cmd_app = cmd_app
+        self.command_color = command_color
+        self.alias_color = alias_color
+        self.macro_color = macro_color
+        self.flag_color = flag_color
+        self.argument_color = argument_color
 
     def lex_document(self, document: Document) -> Callable[[int], Any]:
         """Lex the document."""
@@ -182,16 +207,30 @@ def get_line(lineno: int) -> list[tuple[str, str]]:
 
                 if command:
                     # Determine the style for the command
-                    style = ''
-                    if command in self.cmd_app.get_all_commands():
-                        style = 'ansigreen'
-                    elif command in self.cmd_app.aliases:
-                        style = 'ansicyan'
-                    elif command in self.cmd_app.macros:
-                        style = 'ansimagenta'
-
-                    # Add the command with the determined style
-                    tokens.append((style, command))
+                    shortcut_found = False
+                    for shortcut, _ in self.cmd_app.statement_parser.shortcuts:
+                        if command.startswith(shortcut):
+                            # Add the shortcut with the command style
+                            tokens.append((self.command_color, shortcut))
+
+                            # If there's more in the command word, it's an argument
+                            if len(command) > len(shortcut):
+                                tokens.append((self.argument_color, command[len(shortcut) :]))
+
+                            shortcut_found = True
+                            break
+
+                    if not shortcut_found:
+                        style = ''
+                        if command in self.cmd_app.get_all_commands():
+                            style = self.command_color
+                        elif command in self.cmd_app.aliases:
+                            style = self.alias_color
+                        elif command in self.cmd_app.macros:
+                            style = self.macro_color
+
+                        # Add the command with the determined style
+                        tokens.append((style, command))
 
                 # Add the rest of the line
                 if cmd_end < len(line):
@@ -211,9 +250,9 @@ def get_line(lineno: int) -> list[tuple[str, str]]:
                         if space:
                             tokens.append(('', text))
                         elif flag:
-                            tokens.append(('ansired', text))
+                            tokens.append((self.flag_color, text))
                         elif (quoted or word) and text not in exclude_tokens:
-                            tokens.append(('ansiyellow', text))
+                            tokens.append((self.argument_color, text))
                         else:
                             tokens.append(('', text))
             elif line:

tests/test_pt_utils.py

@@ -25,6 +25,7 @@ def __init__(self):
         self.completion_header = ''
         self.statement_parser = Mock()
         self.statement_parser.terminators = [';']
+        self.statement_parser.shortcuts = []
         self.statement_parser._command_pattern = re.compile(r'\A\s*(\S*?)(\s|\Z)')
         self.aliases = {}
         self.macros = {}
@@ -146,6 +147,25 @@ def test_lex_document_unclosed_quote(self, mock_cmd_app):
 
         assert tokens == [('ansigreen', 'echo'), ('', ' '), ('ansiyellow', '"hello')]
 
+    def test_lex_document_shortcut(self, mock_cmd_app):
+        """Test lexing a shortcut."""
+        mock_cmd_app.statement_parser.shortcuts = [('!', 'shell')]
+        lexer = pt_utils.Cmd2Lexer(cast(Any, mock_cmd_app))
+
+        # Case 1: Shortcut glued to argument
+        line = "!ls"
+        document = Document(line)
+        get_line = lexer.lex_document(document)
+        tokens = get_line(0)
+        assert tokens == [('ansigreen', '!'), ('ansiyellow', 'ls')]
+
+        # Case 2: Shortcut with space
+        line = "! ls"
+        document = Document(line)
+        get_line = lexer.lex_document(document)
+        tokens = get_line(0)
+        assert tokens == [('ansigreen', '!'), ('', ' '), ('ansiyellow', 'ls')]
+
 
 class TestCmd2Completer:
     def test_get_completions_basic(self, mock_cmd_app):