Fixed multiline command highlighting in Cmd2Lexer

Author: tleonhardt

cmd2/pt_utils.py

@@ -219,77 +219,84 @@ def __init__(
 
     def lex_document(self, document: Document) -> Callable[[int], Any]:
         """Lex the document."""
+        # Get redirection tokens and terminators to avoid highlighting them as values
+        exclude_tokens = set(constants.REDIRECTION_TOKENS)
+        exclude_tokens.update(self.cmd_app.statement_parser.terminators)
+        arg_pattern = re.compile(r'(\s+)|(--?[^\s\'"]+)|("[^"]*"?|\'[^\']*\'?)|([^\s\'"]+)')
+
+        def highlight_args(text: str, tokens: list[tuple[str, str]]) -> None:
+            """Highlight arguments in a string."""
+            for m in arg_pattern.finditer(text):
+                space, flag, quoted, word = m.groups()
+                match_text = m.group(0)
+
+                if space:
+                    tokens.append(('', match_text))
+                elif flag:
+                    tokens.append((self.flag_color, match_text))
+                elif (quoted or word) and match_text not in exclude_tokens:
+                    tokens.append((self.argument_color, match_text))
+                else:
+                    tokens.append(('', match_text))
 
         def get_line(lineno: int) -> list[tuple[str, str]]:
             """Return the tokens for the given line number."""
             line = document.lines[lineno]
             tokens: list[tuple[str, str]] = []
 
-            # Use cmd2's command pattern to find the first word (the command)
-            if ru.ALLOW_STYLE != ru.AllowStyle.NEVER and (
-                match := self.cmd_app.statement_parser._command_pattern.search(line)
-            ):
-                # Group 1 is the command, Group 2 is the character(s) that terminated the command match
-                command = match.group(1)
-                cmd_start = match.start(1)
-                cmd_end = match.end(1)
-
-                # Add any leading whitespace
-                if cmd_start > 0:
-                    tokens.append(('', line[:cmd_start]))
-
-                if command:
-                    # Determine the style for the 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):
-                    rest = line[cmd_end:]
-                    # Regex to match whitespace, flags, quoted strings, or other words
-                    arg_pattern = re.compile(r'(\s+)|(--?[^\s\'"]+)|("[^"]*"?|\'[^\']*\'?)|([^\s\'"]+)')
-
-                    # Get redirection tokens and terminators to avoid highlighting them as values
-                    exclude_tokens = set(constants.REDIRECTION_TOKENS)
-                    exclude_tokens.update(self.cmd_app.statement_parser.terminators)
-
-                    for m in arg_pattern.finditer(rest):
-                        space, flag, quoted, word = m.groups()
-                        text = m.group(0)
-
-                        if space:
-                            tokens.append(('', text))
-                        elif flag:
-                            tokens.append((self.flag_color, text))
-                        elif (quoted or word) and text not in exclude_tokens:
-                            tokens.append((self.argument_color, text))
-                        else:
-                            tokens.append(('', text))
-            elif line:
-                # No command match found or colors aren't allowed, add the entire line unstyled
-                tokens.append(('', line))
+            # Only attempt to match a command on the first line
+            if lineno == 0 and ru.ALLOW_STYLE != ru.AllowStyle.NEVER:
+                # Use cmd2's command pattern to find the first word (the command)
+                match = self.cmd_app.statement_parser._command_pattern.search(line)
+                if match:
+                    # Group 1 is the command, Group 2 is the character(s) that terminated the command match
+                    command = match.group(1)
+                    cmd_start = match.start(1)
+                    cmd_end = match.end(1)
+
+                    # Add any leading whitespace
+                    if cmd_start > 0:
+                        tokens.append(('', line[:cmd_start]))
+
+                    if command:
+                        # Determine the style for the 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 as arguments
+                    if cmd_end < len(line):
+                        highlight_args(line[cmd_end:], tokens)
+                else:
+                    # No command match found on the first line
+                    tokens.append(('', line))
+            else:
+                # All other lines are unstyled or treated as arguments
+                if ru.ALLOW_STYLE != ru.AllowStyle.NEVER:
+                    highlight_args(line, tokens)
+                else:
+                    tokens.append(('', line))
 
             return tokens
 

tests/test_pt_utils.py

@@ -210,13 +210,30 @@ def test_lex_document_shortcut(self, mock_cmd_app):
         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')]
 
+    def test_lex_document_multiline(self, mock_cmd_app):
+        """Test lexing a multiline command."""
+        mock_cmd_app.all_commands = ["orate"]
+        lexer = pt_utils.Cmd2Lexer(cast(Any, mock_cmd_app))
+
+        # Command on first line, argument on second line that looks like a command
+        line = "orate\nhelp"
+        document = Document(line)
+        get_line = lexer.lex_document(document)
+
+        # First line should have command
+        tokens0 = get_line(0)
+        assert tokens0 == [('ansigreen', 'orate')]
+
+        # Second line should have argument (not command)
+        tokens1 = get_line(1)
+        assert tokens1 == [('ansiyellow', 'help')]
+
 
 class TestCmd2Completer:
     def test_get_completions(self, mock_cmd_app: MockCmd, monkeypatch) -> None: