Improve python tree-sitter's string fontification

* lisp/progmodes/python.el (python--treesit-fontify-string): Handle
not only f-strings, but also docstrings, and NODE is now the last
quote rather than the whole string.
(python--treesit-settings): Use python--treesit-fontify-string for
every occasion.

1 files changed, 23 insertions, 9 deletions

diff --git a/lisp/progmodes/python.el b/lisp/progmodes/python.el
index 558868efdf7..a9aff167767 100644
--- a/lisp/progmodes/python.el
+++ b/lisp/progmodes/python.el
@@ -1015,12 +1015,24 @@ It makes underscores and dots word constituent chars.")
     "VMSError" "WindowsError"
     ))
 
-(defun python--treesit-fontify-string (beg end _)
-  "Fontify string between BEG and END.
-Do not fontify the initial f for f-strings."
-  (let ((beg (if (eq (char-after beg) ?f)
-                 (1+ beg) beg)))
-    (put-text-property beg end 'face 'font-lock-string-face)))
+(defun python--treesit-fontify-string (_beg _end node)
+  "Fontify string.
+NODE is the last quote in the string.  Do not fontify the initial
+f for f-strings."
+  (let* ((string (treesit-node-parent node))
+         (string-beg (treesit-node-start string))
+         (string-end (treesit-node-end string))
+         (maybe-defun (treesit-node-parent
+                       (treesit-node-parent
+                        (treesit-node-parent string))))
+         (face (if (member (treesit-node-type maybe-defun)
+                           '("function_definition"
+                             "class_definition"))
+                   'font-lock-doc-face
+                 'font-lock-string-face)))
+    (when (eq (char-after string-beg) ?f)
+      (cl-incf string-beg))
+    (put-text-property string-beg string-end 'face face)))
 
 (defvar python--treesit-settings
   (treesit-font-lock-rules
@@ -1031,9 +1043,11 @@ Do not fontify the initial f for f-strings."
    :feature 'string
    :language 'python
    :override t
-   '((string) @python--treesit-fontify-string
-     ((string) @font-lock-doc-face
-      (:match "^\"\"\"" @font-lock-doc-face)))
+   ;; Capture the last quote node rather than the whole string.  The
+   ;; whole string might not be captured if it's not contained in the
+   ;; region being fontified.  E.g., the user inserts a quote, that
+   ;; single quote is the whole region we are fontifying.
+   '((string "\"" @python--treesit-fontify-string :anchor))
 
    :feature 'string-interpolation
    :language 'python