Enhance Python font-lock to support multilines

* test/lisp/progmodes/python-tests.el
(python-tests-assert-faces-after-change): New helper function.
(python-font-lock-keywords-level-1-3)
(python-font-lock-assignment-statement-multiline-*): New tests.

* lisp/progmodes/python.el (python-rx): Add `sp-nl' to represent
space or newline (with/without backslash).
(python-font-lock-keywords-level-1)
(python-font-lock-keywords-maximum-decoration): Allow newlines
where appropriate.
(python-font-lock-extend-region): New function.
(python-mode): Set `python-font-lock-extend-region' to
`font-lock-extend-after-change-region-function'.

2 files changed, 145 insertions, 16 deletions

diff --git a/lisp/progmodes/python.el b/lisp/progmodes/python.el
index 44df3186b27..e1350391994 100644
--- a/lisp/progmodes/python.el
+++ b/lisp/progmodes/python.el
@@ -359,6 +359,7 @@
   "Python mode specialized rx macro.
 This variant of `rx' supports common Python named REGEXPS."
   `(rx-let ((sp-bsnl (or space (and ?\\ ?\n)))
+            (sp-nl (or space (and (? ?\\) ?\n)))
             (block-start       (seq symbol-start
                                     (or "def" "class" "if" "elif" "else" "try"
                                         "except" "finally" "for" "while" "with"
@@ -583,9 +584,9 @@ the {...} holes that appear within f-strings."
              finally return (and result-valid result))))
 
 (defvar python-font-lock-keywords-level-1
-  `((,(python-rx symbol-start "def" (1+ space) (group symbol-name))
+  `((,(python-rx symbol-start "def" (1+ sp-bsnl) (group symbol-name))
      (1 font-lock-function-name-face))
-    (,(python-rx symbol-start "class" (1+ space) (group symbol-name))
+    (,(python-rx symbol-start "class" (1+ sp-bsnl) (group symbol-name))
      (1 font-lock-type-face)))
   "Font lock keywords to use in `python-mode' for level 1 decoration.
 
@@ -725,12 +726,12 @@ sign in chained assignment."
     ;;   [*a] = 5, 6
     ;; are handled separately below
     (,(python-font-lock-assignment-matcher
-        (python-rx (? (or "[" "(") (* space))
-                   grouped-assignment-target (* space) ?, (* space)
-                   (* assignment-target (* space) ?, (* space))
-                   (? assignment-target (* space))
-                   (? ?, (* space))
-                   (? (or ")" "]") (* space))
+        (python-rx (? (or "[" "(") (* sp-nl))
+                   grouped-assignment-target (* sp-nl) ?, (* sp-nl)
+                   (* assignment-target (* sp-nl) ?, (* sp-nl))
+                   (? assignment-target (* sp-nl))
+                   (? ?, (* sp-nl))
+                   (? (or ")" "]") (* sp-bsnl))
                    (group assignment-operator)))
      (1 font-lock-variable-name-face)
      (,(python-rx grouped-assignment-target)
@@ -745,19 +746,20 @@ sign in chained assignment."
     ;;   c: Collection = {1, 2, 3}
     ;;   d: Mapping[int, str] = {1: 'bar', 2: 'baz'}
     (,(python-font-lock-assignment-matcher
-        (python-rx grouped-assignment-target (* space)
-                   (? ?: (* space) (+ not-simple-operator) (* space))
-                   assignment-operator))
+       (python-rx (or line-start ?\;) (* sp-bsnl)
+                  grouped-assignment-target (* sp-bsnl)
+                  (? ?: (* sp-bsnl) (+ not-simple-operator) (* sp-bsnl))
+                  assignment-operator))
      (1 font-lock-variable-name-face))
     ;; special cases
     ;;   (a) = 5
     ;;   [a] = 5,
     ;;   [*a] = 5, 6
     (,(python-font-lock-assignment-matcher
-       (python-rx (or line-start ?\; ?=) (* space)
-                  (or "[" "(") (* space)
-                  grouped-assignment-target (* space)
-                  (or ")" "]") (* space)
+       (python-rx (or line-start ?\; ?=) (* sp-bsnl)
+                  (or "[" "(") (* sp-nl)
+                  grouped-assignment-target (* sp-nl)
+                  (or ")" "]") (* sp-bsnl)
                   assignment-operator))
      (1 font-lock-variable-name-face))
     ;; escape sequences within bytes literals
@@ -796,6 +798,18 @@ decorators, exceptions, and assignments.")
 Which one will be chosen depends on the value of
 `font-lock-maximum-decoration'.")
 
+(defun python-font-lock-extend-region (beg end _old-len)
+  "Extend font-lock region given by BEG and END to statement boundaries."
+  (save-excursion
+    (save-match-data
+      (goto-char beg)
+      (python-nav-beginning-of-statement)
+      (setq beg (point))
+      (goto-char end)
+      (python-nav-end-of-statement)
+      (setq end (point))
+      (cons beg end))))
+
 
 (defconst python-syntax-propertize-function
   (syntax-propertize-rules
@@ -5780,7 +5794,9 @@ REPORT-FN is Flymake's callback function."
               `(,python-font-lock-keywords
                 nil nil nil nil
                 (font-lock-syntactic-face-function
-                 . python-font-lock-syntactic-face-function)))
+                 . python-font-lock-syntactic-face-function)
+                (font-lock-extend-after-change-region-function
+                 . python-font-lock-extend-region)))
 
   (setq-local syntax-propertize-function
               python-syntax-propertize-function)
diff --git a/test/lisp/progmodes/python-tests.el b/test/lisp/progmodes/python-tests.el
index 9e8fa7f5520..875c92573ef 100644
--- a/test/lisp/progmodes/python-tests.el
+++ b/test/lisp/progmodes/python-tests.el
@@ -108,6 +108,20 @@ STRING, it is skipped so the next STRING occurrence is selected."
            while pos
            collect (cons pos (get-text-property pos 'face))))
 
+(defun python-tests-assert-faces-after-change (content faces search replace)
+  "Assert that font faces for CONTENT are equal to FACES after change.
+All occurrences of SEARCH are changed to REPLACE."
+  (python-tests-with-temp-buffer
+   content
+   ;; Force enable font-lock mode without jit-lock.
+   (rename-buffer "*python-font-lock-test*" t)
+   (let (noninteractive font-lock-support-mode)
+     (font-lock-mode))
+   (while
+       (re-search-forward search nil t)
+     (replace-match replace))
+   (should (equal faces (python-tests-get-buffer-faces)))))
+
 (defun python-tests-self-insert (char-or-str)
   "Call `self-insert-command' for chars in CHAR-OR-STR."
   (let ((chars
@@ -226,6 +240,13 @@ aliqua."
    "def 1func():"
    '((1 . font-lock-keyword-face) (4))))
 
+(ert-deftest python-font-lock-keywords-level-1-3 ()
+  (python-tests-assert-faces
+   "def \\
+        func():"
+   '((1 . font-lock-keyword-face) (4)
+     (15 . font-lock-function-name-face) (19))))
+
 (ert-deftest python-font-lock-assignment-statement-1 ()
   (python-tests-assert-faces
    "a, b, c = 1, 2, 3"
@@ -380,6 +401,98 @@ def f(x: CustomInt) -> CustomInt:
      (128 . font-lock-builtin-face) (131)
      (144 . font-lock-keyword-face) (150))))
 
+(ert-deftest python-font-lock-assignment-statement-multiline-1 ()
+  (python-tests-assert-faces-after-change
+   "
+[
+    a,
+    b
+] # (
+    1,
+    2
+)
+"
+   '((1)
+     (8 . font-lock-variable-name-face) (9)
+     (15 . font-lock-variable-name-face) (16))
+   "#" "="))
+
+(ert-deftest python-font-lock-assignment-statement-multiline-2 ()
+  (python-tests-assert-faces-after-change
+   "
+[
+    *a
+] # 5, 6
+"
+   '((1)
+     (9 . font-lock-variable-name-face) (10))
+   "#" "="))
+
+(ert-deftest python-font-lock-assignment-statement-multiline-3 ()
+  (python-tests-assert-faces-after-change
+   "a\\
+    ,\\
+    b\\
+    ,\\
+    c\\
+    #\\
+    1\\
+    ,\\
+    2\\
+    ,\\
+    3"
+   '((1 . font-lock-variable-name-face) (2)
+     (15 . font-lock-variable-name-face) (16)
+     (29 . font-lock-variable-name-face) (30))
+   "#" "="))
+
+(ert-deftest python-font-lock-assignment-statement-multiline-4 ()
+  (python-tests-assert-faces-after-change
+   "a\\
+    :\\
+    int\\
+    #\\
+    5"
+   '((1 . font-lock-variable-name-face) (2)
+     (15 . font-lock-builtin-face) (18))
+   "#" "="))
+
+(ert-deftest python-font-lock-assignment-statement-multiline-5 ()
+  (python-tests-assert-faces-after-change
+   "(\\
+    a\\
+)\\
+    #\\
+    5\\
+    ;\\
+    (\\
+    b\\
+    )\\
+    #\\
+    6"
+   '((1)
+     (8 . font-lock-variable-name-face) (9)
+     (46 . font-lock-variable-name-face) (47))
+   "#" "="))
+
+(ert-deftest python-font-lock-assignment-statement-multiline-6 ()
+  (python-tests-assert-faces-after-change
+   "(
+    a
+)\\
+    #\\
+    5\\
+    ;\\
+    (
+    b
+    )\\
+    #\\
+    6"
+   '((1)
+     (7 . font-lock-variable-name-face) (8)
+     (43 . font-lock-variable-name-face) (44))
+   "#" "="))
+
 (ert-deftest python-font-lock-escape-sequence-string-newline ()
   (python-tests-assert-faces
    "'\\n'