Fix more cases where lazy is not highlighted in repl (#36)

Author: StanFromIreland

Lib/_pyrepl/utils.py

@@ -276,7 +276,11 @@ def is_soft_keyword_used(*tokens: TI | None) -> bool:
             TI(T.NAME, string=s)
         ):
             return not keyword.iskeyword(s)
-        case (None | TI(T.NEWLINE) | TI(T.INDENT) | TI(T.DEDENT), TI(string="lazy"), TI(string="import") | TI(string="from")):
+        case (
+            None | TI(T.NEWLINE) | TI(T.INDENT) | TI(T.DEDENT) | TI(string=":" | ";"),
+            TI(string="lazy"),
+            TI(string="import") | TI(string="from")
+        ):
             return True
         case _:
             return False

Lib/test/test_pyrepl/test_utils.py

@@ -96,7 +96,9 @@ def test_gen_colors_keyword_highlighting(self):
             ("list", [("list", "builtin")]),
             ("    \n dict", [("dict", "builtin")]),
             ("    lazy import", [("lazy", "soft_keyword"), ("import", "keyword")]),
-            ("lazy from cool_people import pablo", [('lazy', 'soft_keyword'), ('from', 'keyword'), ('import', 'keyword')])
+            ("lazy from cool_people import pablo", [('lazy', 'soft_keyword'), ('from', 'keyword'), ('import', 'keyword')]),
+            ("if sad: lazy import happy", [("if", "keyword"), (":", "op"), ("lazy", "soft_keyword"), ("import", "keyword")]),
+            ("pass; lazy import z", [("pass", "keyword"), (';', 'op'), ("lazy", "soft_keyword"), ("import", "keyword")])
         ]
         for code, expected_highlights in cases:
             with self.subTest(code=code):