bug#67061: [PATCH] Improve syntax highlighting for python-ts-mode

[Denis Zubarev <dvzubarev@yandex.ru>]

[-- Attachment #1: Type: text/plain, Size: 1517 bytes --]

Tags: patch


Improve syntax highlighting for python-ts-mode.

- Fontify compound keywords "not in" and "is not" (fixes bug#67015)
- Fix fontification of strings inside of f-strings interpolation,
   e.g. for f"beg {'nested'}" - 'nested' was not fontified as string.
- Do not override the face of builtin functions (all, bytes etc.) with
  the function call face
- Add missing assignment expressions
- Fontify built-ins (dict,list,etc.) as types when they are used in type hints

I have added tests to treesit-tests.el, but I'm not sure that it is the
right place for them. I tried to create new file
python-ts-mode-tests.el, but when I run tests using make, I got
an error:
*** No rule to make target '../lisp/progmodes/python-ts-mode.el', needed by 'lisp/progmodes/python-ts-mode-tests.log'.  Stop.

It seems that python-ts-mode is required to be in its own file, but it
is defined in python.el.
Adding tests to python-tests.el seems like not good idea too, because
some CI code is searching for *-ts-mode-tests.el or treesit-tests.el files.


In GNU Emacs 30.0.50 (build 3, x86_64-pc-linux-gnu, GTK+ Version
 3.24.33, cairo version 1.16.0) of 2023-11-11 built on NUC-here
Repository revision: 400a71b8f2c5a49dce4f542adfd2fdb59eb34243
Repository branch: master
Windowing system distributor 'The X.Org Foundation', version 11.0.12101004
System Description: Ubuntu 22.04.3 LTS

Configured using:
 'configure --with-modules --with-native-compilation=aot
 --with-imagemagick --with-json --with-tree-sitter --with-xft'


[-- Warning: decoded text below may be mangled, UTF-8 assumed --]
[-- Attachment #2: 0001-Improve-syntax-highlighting-for-python-ts-mode.patch --]
[-- Type: text/patch, Size: 11271 bytes --]

From 0b3165eb90bdc1b85688d3c9a1902c1f293972d1 Mon Sep 17 00:00:00 2001
From: Denis Zubarev <dvzubarev@yandex.ru>
Date: Sat, 11 Nov 2023 04:55:44 +0300
Subject: [PATCH] Improve syntax highlighting for python-ts-mode

* lisp/progmodes/python.el (python--treesit-keywords): Fontify compound
keywords "not in" and "is not" (bug#67015).
(python--treesit-fontify-string): Fix fontification of strings inside of
f-strings interpolation, e.g. for f"beg {'nested'}" - 'nested' was not
fontified as string.
(python--treesit-fontify-string-interpolation): Remove function.
(python--treesit-settings): Do not override the face of builtin
functions (all, bytes etc.) with the function call face.  Add missing
assignment expressions.  Fontify built-ins (dict,list,etc.) as types
when they are used in type hints.
* test/src/treesit-tests.el (python-ts-tests-with-temp-buffer):
(python-ts-mode-compound-keywords-face):
(python-ts-mode-named-assignement-face):
(python-ts-mode-nested-types-face):
(python-ts-mode-builtin-call-face):
(python-ts-mode-interpolation-nested-string):
(python-ts-mode-interpolation-doc-string): Add tests.
---
 lisp/progmodes/python.el  | 81 +++++++++++++++++-----------------
 test/src/treesit-tests.el | 93 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 134 insertions(+), 40 deletions(-)

diff --git a/lisp/progmodes/python.el b/lisp/progmodes/python.el
index d3cb5a77e2..e3a4e0d2f9 100644
--- a/lisp/progmodes/python.el
+++ b/lisp/progmodes/python.el
@@ -979,7 +979,7 @@ python--treesit-keywords
     "raise" "return" "try" "while" "with" "yield"
     ;; These are technically operators, but we fontify them as
     ;; keywords.
-    "and" "in" "is" "not" "or"))
+    "and" "in" "is" "not" "or" "is not" "not in"))
 
 (defvar python--treesit-builtins
   '("abs" "all" "any" "ascii" "bin" "bool" "breakpoint" "bytearray"
@@ -1042,9 +1042,7 @@ python--treesit-fontify-string
 f-strings.  OVERRIDE is the override flag described in
 `treesit-font-lock-rules'.  START and END mark the region to be
 fontified."
-  (let* ((string-beg (treesit-node-start node))
-         (string-end (treesit-node-end node))
-         (maybe-expression (treesit-node-parent node))
+  (let* ((maybe-expression (treesit-node-parent node))
          (grandparent (treesit-node-parent
                        (treesit-node-parent
                         maybe-expression)))
@@ -1072,28 +1070,29 @@ python--treesit-fontify-string
                         (equal (treesit-node-type maybe-expression)
                                "expression_statement"))
                    'font-lock-doc-face
-                 'font-lock-string-face)))
-    ;; Don't highlight string prefixes like f/r/b.
-    (save-excursion
-      (goto-char string-beg)
-      (when (re-search-forward "[\"']" string-end t)
-        (setq string-beg (match-beginning 0))))
-    (treesit-fontify-with-override
-     string-beg string-end face override start end)))
+                 'font-lock-string-face))
+         ;; We will highlight only string_start/string_content/string_end children
+         ;; Do not touch interpolation node that can occur inside of the string
+         (string-nodes (treesit-filter-child
+                        node
+                        (lambda (ch) (member (treesit-node-type ch)
+                                             '("string_start"
+                                               "string_content"
+                                               "string_end")))
+                        t)))
+
+    (dolist (string-node string-nodes)
+      (let ((string-beg (treesit-node-start string-node))
+            (string-end (treesit-node-end string-node)))
+        (when (equal (treesit-node-type string-node) "string_start")
+          ;; Don't highlight string prefixes like f/r/b.
+          (save-excursion
+            (goto-char string-beg)
+            (when (re-search-forward "[\"']" string-end t)
+              (setq string-beg (match-beginning 0)))))
 
-(defun python--treesit-fontify-string-interpolation
-    (node _ start end &rest _)
-  "Fontify string interpolation.
-NODE is the string node.  Do not fontify the initial f for
-f-strings.  START and END mark the region to be
-fontified."
-  ;; This is kind of a hack, it basically removes the face applied by
-  ;; the string feature, so that following features can apply their
-  ;; face.
-  (let ((n-start (treesit-node-start node))
-        (n-end (treesit-node-end node)))
-    (remove-text-properties
-     (max start n-start) (min end n-end) '(face))))
+        (treesit-fontify-with-override
+         string-beg string-end face override start end)))))
 
 (defvar python--treesit-settings
   (treesit-font-lock-rules
@@ -1103,14 +1102,9 @@ python--treesit-settings
 
    :feature 'string
    :language 'python
-   '((string) @python--treesit-fontify-string)
+   '((string) @python--treesit-fontify-string
+     (interpolation ["{" "}"] @font-lock-misc-punctuation-face))
 
-   ;; HACK: This feature must come after the string feature and before
-   ;; other features.  Maybe we should make string-interpolation an
-   ;; option rather than a feature.
-   :feature 'string-interpolation
-   :language 'python
-   '((interpolation) @python--treesit-fontify-string-interpolation)
 
    :feature 'keyword
    :language 'python
@@ -1126,12 +1120,6 @@ python--treesit-settings
       name: (identifier) @font-lock-type-face)
      (parameters (identifier) @font-lock-variable-name-face))
 
-   :feature 'function
-   :language 'python
-   '((call function: (identifier) @font-lock-function-call-face)
-     (call function: (attribute
-                      attribute: (identifier) @font-lock-function-call-face)))
-
    :feature 'builtin
    :language 'python
    `(((identifier) @font-lock-builtin-face
@@ -1142,6 +1130,12 @@ python--treesit-settings
                       eol))
               @font-lock-builtin-face)))
 
+   :feature 'function
+   :language 'python
+   '((call function: (identifier) @font-lock-function-call-face)
+     (call function: (attribute
+                      attribute: (identifier) @font-lock-function-call-face)))
+
    :feature 'constant
    :language 'python
    '([(true) (false) (none)] @font-lock-constant-face)
@@ -1154,6 +1148,10 @@ python--treesit-settings
      (assignment left: (attribute
                         attribute: (identifier)
                         @font-lock-property-use-face))
+     (augmented_assignment left: (identifier)
+                           @font-lock-variable-name-face)
+     (named_expression name: (identifier)
+                       @font-lock-variable-name-face)
      (pattern_list (identifier)
                    @font-lock-variable-name-face)
      (tuple_pattern (identifier)
@@ -1168,15 +1166,18 @@ python--treesit-settings
    '((decorator "@" @font-lock-type-face)
      (decorator (call function: (identifier) @font-lock-type-face))
      (decorator (identifier) @font-lock-type-face))
-
    :feature 'type
    :language 'python
+   ;; override built-in faces when dict/list are used for type hints
+   :override t
    `(((identifier) @font-lock-type-face
       (:match ,(rx-to-string
                 `(seq bol (or ,@python--treesit-exceptions)
                       eol))
               @font-lock-type-face))
-     (type (identifier) @font-lock-type-face))
+     (type [(identifier) (none)] @font-lock-type-face)
+     (type
+      (_ [(identifier) (none)] @font-lock-type-face)))
 
    :feature 'escape-sequence
    :language 'python
diff --git a/test/src/treesit-tests.el b/test/src/treesit-tests.el
index 791e902bd0..71c7b92957 100644
--- a/test/src/treesit-tests.el
+++ b/test/src/treesit-tests.el
@@ -1167,6 +1167,99 @@ treesit-defun-navigation-top-level
    treesit--ert-defun-navigation-top-level-master
    'top-level))
 
+;;; python-ts-mode font-lock tests
+
+(defmacro python-ts-tests-with-temp-buffer (contents &rest body)
+  "Create a `python-ts-mode' enabled temp buffer with CONTENTS.
+BODY is code to be executed within the temp buffer.  Point is
+always located at the beginning of buffer."
+  (declare (indent 1) (debug t))
+  `(with-temp-buffer
+     (skip-unless (treesit-language-available-p 'python))
+     (require 'python)
+     (let ((python-indent-guess-indent-offset nil))
+       (python-ts-mode)
+       (insert ,contents)
+       (goto-char (point-min))
+       ,@body)))
+
+(ert-deftest python-ts-mode-compound-keywords-face ()
+  (dolist (test '("is not" "not in"))
+    (python-ts-tests-with-temp-buffer (concat "t " test " t")
+      (font-lock-ensure)
+      (forward-to-word)
+      (should (eq (face-at-point) font-lock-keyword-face))
+      (forward-to-word)
+      (should (eq (face-at-point) font-lock-keyword-face)))))
+
+(ert-deftest python-ts-mode-named-assignement-face ()
+  (python-ts-tests-with-temp-buffer "var := 3"
+    (font-lock-ensure)
+    (should (eq (face-at-point) font-lock-variable-name-face))))
+
+(ert-deftest python-ts-mode-nested-types-face ()
+  (python-ts-tests-with-temp-buffer "def func(v:dict[ list[ tuple[str] ], int | None] | None):"
+    (font-lock-ensure)
+    (dolist (test '("dict" "list" "tuple" "str" "int" "None" "None"))
+      (search-forward test)
+      (goto-char (match-beginning 0))
+      (should (eq (face-at-point) font-lock-type-face)))))
+
+(ert-deftest python-ts-mode-builtin-call-face ()
+  (python-ts-tests-with-temp-buffer "all()"
+    ;; enable 'function' feature from 4th level
+    (setopt treesit-font-lock-level 4)
+    (font-lock-ensure)
+    (should (eq (face-at-point) font-lock-builtin-face))))
+
+(ert-deftest python-ts-mode-interpolation-nested-string ()
+  (python-ts-tests-with-temp-buffer "t = f\"beg {var + 'string'}\""
+    (font-lock-ensure)
+
+    (search-forward "var")
+    (goto-char (match-beginning 0))
+    (should (eq (face-at-point) font-lock-variable-name-face))
+
+    (goto-char (point-min))
+    (dolist (test '("f" "{" "+" "}"))
+      (search-forward test)
+      (goto-char (match-beginning 0))
+      (should (not (eq (face-at-point) font-lock-string-face))))
+
+
+    (goto-char (point-min))
+    (dolist (test '("beg" "'string'" "\""))
+      (search-forward test)
+      (goto-char (match-beginning 0))
+      (should (eq (face-at-point) font-lock-string-face)))))
+
+(ert-deftest python-ts-mode-interpolation-doc-string ()
+  (python-ts-tests-with-temp-buffer "f\"\"\"beg {'s1' + var + 's2'} end\"\"\""
+    (font-lock-ensure)
+
+    (search-forward "var")
+    (goto-char (match-beginning 0))
+    (should (eq (face-at-point) font-lock-variable-name-face))
+
+    (goto-char (point-min))
+    (dolist (test '("f" "{" "+" "}"))
+      (search-forward test)
+      (goto-char (match-beginning 0))
+      (should (not (eq (face-at-point) font-lock-string-face))))
+
+    (goto-char (point-min))
+    (dolist (test '("\"\"\"" "beg" "end" "\"\"\""))
+      (search-forward test)
+      (goto-char (match-beginning 0))
+      (should (eq (face-at-point) font-lock-doc-face)))
+
+    (goto-char (point-min))
+    (dolist (test '("'s1'" "'s2'"))
+      (search-forward test)
+      (goto-char (match-beginning 0))
+      (should (eq (face-at-point) font-lock-string-face)))))
+
+
 ;; TODO
 ;; - Functions in treesit.el
 ;; - treesit-load-name-override-list
-- 
2.34.1