Improve fontification of Python assignments with type hints

* lisp/progmodes/python.el (python-font-lock-keywords-maximum-decoration):
Fontify type hints of assignment statement.  (Bug#69357)
* test/lisp/progmodes/python-tests.el
(python-font-lock-assignment-statement-11)
(python-font-lock-assignment-statement-12)
(python-font-lock-assignment-statement-13)
(python-font-lock-assignment-statement-18): Add fontification of type
hints.
(python-font-lock-assignment-statement-19): New test.

2 files changed, 39 insertions, 12 deletions

diff --git a/lisp/progmodes/python.el b/lisp/progmodes/python.el
index 831bec7f4af..360936c9efd 100644
--- a/lisp/progmodes/python.el
+++ b/lisp/progmodes/python.el
@@ -794,6 +794,25 @@ sign in chained assignment."
            )
           symbol-end)
      . font-lock-type-face)
+    ;; single assignment with/without type hints, e.g.
+    ;;   a: int = 5
+    ;;   b: Tuple[Optional[int], Union[Sequence[str], str]] = (None, 'foo')
+    ;;   c: Collection = {1, 2, 3}
+    ;;   d: Mapping[int, str] = {1: 'bar', 2: 'baz'}
+    (,(python-font-lock-assignment-matcher
+       (python-rx (or line-start ?\;) (* space)
+                  grouped-assignment-target (* space)
+                  (? ?: (* space) (group (+ not-simple-operator)) (* space))
+                  (group assignment-operator)))
+     (1 font-lock-variable-name-face)
+     (3 'font-lock-operator-face)
+     (,(python-rx symbol-name)
+      (progn
+        (when-let ((type-start (match-beginning 2)))
+          (goto-char type-start))
+        (match-end 0))
+      nil
+      (0 font-lock-type-face)))
     ;; multiple assignment
     ;; (note that type hints are not allowed for multiple assignments)
     ;;   a, b, c = 1, 2, 3
@@ -826,18 +845,6 @@ sign in chained assignment."
         (match-beginning 2))            ; limit the search until the assignment
       nil
       (1 font-lock-variable-name-face)))
-    ;; single assignment with type hints, e.g.
-    ;;   a: int = 5
-    ;;   b: Tuple[Optional[int], Union[Sequence[str], str]] = (None, 'foo')
-    ;;   c: Collection = {1, 2, 3}
-    ;;   d: Mapping[int, str] = {1: 'bar', 2: 'baz'}
-    (,(python-font-lock-assignment-matcher
-       (python-rx (or line-start ?\;) (* space)
-                  grouped-assignment-target (* space)
-                  (? ?: (* space) (+ not-simple-operator) (* space))
-                  (group assignment-operator)))
-     (1 font-lock-variable-name-face)
-     (2 'font-lock-operator-face))
     ;; special cases
     ;;   (a) = 5
     ;;   [a] = 5,
diff --git a/test/lisp/progmodes/python-tests.el b/test/lisp/progmodes/python-tests.el
index f50797953c3..2960cca2c06 100644
--- a/test/lisp/progmodes/python-tests.el
+++ b/test/lisp/progmodes/python-tests.el
@@ -391,7 +391,11 @@ p = (1 + 2)
   (python-tests-assert-faces
    "b: Tuple[Optional[int], Union[Sequence[str], str]] = (None, 'foo')"
    '((1 . font-lock-variable-name-face) (2)
+     (4 . font-lock-type-face) (9)
+     (10 . font-lock-type-face) (18)
      (19 . font-lock-builtin-face) (22)
+     (25 . font-lock-type-face) (30)
+     (31 . font-lock-type-face) (39)
      (40 . font-lock-builtin-face) (43)
      (46 . font-lock-builtin-face) (49)
      (52 . font-lock-operator-face) (53)
@@ -402,12 +406,14 @@ p = (1 + 2)
   (python-tests-assert-faces
    "c: Collection = {1, 2, 3}"
    '((1 . font-lock-variable-name-face) (2)
+     (4 . font-lock-type-face) (14)
      (15 . font-lock-operator-face) (16))))
 
 (ert-deftest python-font-lock-assignment-statement-13 ()
   (python-tests-assert-faces
    "d: Mapping[int, str] = {1: 'bar', 2: 'baz'}"
    '((1 . font-lock-variable-name-face) (2)
+     (4 . font-lock-type-face) (11)
      (12 . font-lock-builtin-face) (15)
      (17 . font-lock-builtin-face) (20)
      (22 . font-lock-operator-face) (23)
@@ -472,14 +478,28 @@ def f(x: CustomInt) -> CustomInt:
      (58 . font-lock-operator-face) (59)
      (62 . font-lock-operator-face) (63)
      (70 . font-lock-variable-name-face) (72)
+     (74 . font-lock-type-face) (82)
+     (83 . font-lock-type-face) (92)
      (94 . font-lock-operator-face) (95)
      (102 . font-lock-operator-face) (103)
      (111 . font-lock-variable-name-face) (114)
+     (116 . font-lock-type-face) (125)
      (126 . font-lock-operator-face) (127)
      (128 . font-lock-builtin-face) (131)
      (136 . font-lock-operator-face) (137)
      (144 . font-lock-keyword-face) (150))))
 
+(ert-deftest python-font-lock-assignment-statement-19 ()
+  (python-tests-assert-faces
+   "a: List[List[CustomInt], List[CustomInt]] = []"
+   '((1 . font-lock-variable-name-face) (2)
+     (4 . font-lock-type-face) (8)
+     (9 . font-lock-type-face) (13)
+     (14 . font-lock-type-face) (23)
+     (26 . font-lock-type-face) (30)
+     (31 . font-lock-type-face) (40)
+     (43 . font-lock-operator-face) (44))))
+
 (ert-deftest python-font-lock-operator-1 ()
   (python-tests-assert-faces
    "1 << 2 ** 3 == +4%-5|~6&7^8%9"