Colorize commands separately in the REPL

Author: johnslavik

Lib/_colorize.py

@@ -407,6 +407,7 @@ class Syntax(ThemeSection):
     keyword: str = ANSIColors.BOLD_BLUE
     keyword_constant: str = ANSIColors.BOLD_BLUE
     builtin: str = ANSIColors.CYAN
+    command: str = ANSIColors.BOLD_CYAN
     comment: str = ANSIColors.RED
     string: str = ANSIColors.GREEN
     number: str = ANSIColors.YELLOW

Lib/_pyrepl/utils.py

@@ -24,6 +24,7 @@
 IDENTIFIERS_AFTER = frozenset({"def", "class"})
 KEYWORD_CONSTANTS = frozenset({"True", "False", "None"})
 BUILTINS = frozenset({str(name) for name in dir(builtins) if not name.startswith('_')})
+COMMANDS = frozenset({"exit", "quit", "copyright", "help", "clear"})
 
 
 def THEME(**kwargs):
@@ -235,6 +236,13 @@ def gen_colors_from_token_stream(
                 ):
                     span = Span.from_token(token, line_lengths)
                     yield ColorSpan(span, "soft_keyword")
+                elif (
+                    token.string in COMMANDS
+                    and not prev_token
+                    and next_token.type == T.NEWLINE
+                ):
+                    span = Span.from_token(token, line_lengths)
+                    yield ColorSpan(span, "command")
                 elif (
                     token.string in BUILTINS
                     and not (prev_token and prev_token.exact_type == T.DOT)

Lib/test/test_pyrepl/test_utils.py

@@ -159,3 +159,28 @@ def test_gen_colors_keyword_highlighting(self):
                     span_text = code[color.span.start:color.span.end + 1]
                     actual_highlights.append((span_text, color.tag))
                 self.assertEqual(actual_highlights, expected_highlights)
+
+    def test_gen_colors_command_highlighting(self):
+        cases = [
+            # highlights bare command names
+            ("exit", [("exit", "command")]),
+            ("quit", [("quit", "command")]),
+            ("copyright", [("copyright", "command")]),
+            ("help", [("help", "command")]),
+            ("clear", [("clear", "command")]),
+            # no highlight when not the only token on the line
+            ("x = exit", [("=", "op"), ("exit", "builtin")]),
+            ("obj.exit", [(".", "op")]),
+            # falls through to builtin when called as function or used in expression
+            ("exit()", [("exit", "builtin"), ("(", "op"), (")", "op")]),
+            ("quit(0)", [("quit", "builtin"), ("(", "op"), ("0", "number"), (")", "op")]),
+            ("print(exit)", [("print", "builtin"), ("(", "op"), ("exit", "builtin"), (")", "op")]),
+        ]
+        for code, expected_highlights in cases:
+            with self.subTest(code=code):
+                colors = list(gen_colors(code))
+                actual_highlights = []
+                for color in colors:
+                    span_text = code[color.span.start:color.span.end + 1]
+                    actual_highlights.append((span_text, color.tag))
+                self.assertEqual(actual_highlights, expected_highlights)