(Hopefully) Finished zsh-ts-mode.el

2026-04-14 00:11:04 -07:00
parent 6d69322351
commit be733cb11a
2 changed files with 109 additions and 37 deletions
--- a/elisp/zsh-ts-mode.el
+++ b/elisp/zsh-ts-mode.el
@@ -8,6 +8,11 @@
 (require 'treesit)
 (require 'rx)

+(defgroup zsh-ts-mode nil
+  "Tree-sitter powered Zsh editing."
+  :group 'sh
+  :prefix "zsh-ts-mode-")
+
 ;;; #################
 ;;; # Fontification #
 ;;; #################
@@ -43,15 +48,26 @@ Fontify NODE, an argument to a Zsh \"let\" statement."
                                        'font-lock-variable-name-face))))))

 (defvar zsh-ts-mode--arithmetic-variable-query
-  (treesit-query-compile 'zsh '((word) @name))
-  "Query used in `zsh-ts-mode--fontify-arithmetic-variables'.")
+  (treesit-query-compile
+   'zsh '(((word) @use)
+          (binary_expression left: (variable_name) @assign
+                             operator: "=")))
+  "Query used in `zsh-ts-mode--arithmatic-variables'.")
+
+(defun zsh-ts-mode--arithmatic-variables (parent also-define)
+  "Return a list of arithmetic variables in the arithmetic_expansion PARENT.
+If ALSO-DEFINE is non-nil, also capture defined (e.g. x=2)."
+  (cl-loop for entry in (treesit-query-capture
+                         parent zsh-ts-mode--arithmetic-variable-query)
+           when (or also-define
+                    (eq (car entry) 'use))
+           collect (cdr entry)))

 (defun zsh-ts-mode--fontify-arithmetic-variables
    (node _override _start _end &rest _r)
  "Fontify all arithmetic variables below NODE.
 NODE should be an `arithmetic_expansion' node."
-  (dolist (child (treesit-query-capture
-                  node zsh-ts-mode--arithmetic-variable-query nil nil t))
+  (dolist (child (zsh-ts-mode--arithmatic-variables node nil))
    (let ((start (treesit-node-start child))
          (end (treesit-node-end child)))
      (font-lock-append-text-property start end 'face
@@ -104,13 +120,8 @@ command.")
   :feature 'variable-reference
   :language 'zsh
   :override t
-   '((expansion (simple_variable_name) @font-lock-variable-use-face)
+   '((simple_variable_name) @font-lock-variable-use-face
     (special_variable_name) @font-lock-variable-use-face
-     (expansion_with_modifier (simple_variable_name)
-                              @font-lock-variable-use-face)
-     (expansion_pattern (simple_variable_name) @font-lock-variable-use-face)
-     (expansion_substring (simple_variable_name) @font-lock-variable-use-face)
-     (variable_ref (simple_variable_name) @font-lock-variable-use-face)
     ;; variables in math expansions
     ((arithmetic_expansion :anchor "$((")
      @zsh-ts-mode--fontify-arithmetic-variables))
@@ -269,28 +280,41 @@ Return non-nil if NODE is not an anonymous function."
           name: ([(simple_variable_name)
                   (special_variable_name)]
                  @var))
+          (expansion_default
+           name: ([(simple_variable_name)
+                   (special_variable_name)]
+                  @var))
+          ((variable_name) @var)
          ((word) @func
           (:pred zsh-ts-mode--not-let-predicate @func))))
  "Query for `zsh-ts-mode' `xref-backend-identifier-at-point'.")

+(cl-defun zsh-ts-mode--extract-variable-from-face-at (&optional (pos (point)))
+  "Extract the variable under POS from its face."
+  (cl-flet ((var-face-p (pos)
+              (memq 'font-lock-variable-name-face
+                    (ensure-list (get-text-property pos 'face)))))
+    (when (var-face-p pos)
+      (let ((start (if (and (not (bobp)) (var-face-p (1- pos)))
+                       (or (previous-single-property-change pos 'face)
+                           (point-min))
+                     (point)))
+            (end (or (next-single-property-change (point) 'face)
+                     (point-max))))
+        (propertize
+         (buffer-substring-no-properties start end) 'type 'variable)))))
+
 (cl-defmethod xref-backend-identifier-at-point ((_backend (eql zsh-ts)))
  "`zsh-ts-mode' implementation for xref's identifier at point function."

-  (or (when (memq 'font-lock-variable-name-face
-                  (ensure-list (get-text-property (point) 'face)))
-        (let ((start (previous-single-property-change (point) 'face))
-              (end (next-single-property-change (point) 'face)))
-          (propertize
-           (buffer-substring-no-properties (if start (1+ start) (point-min))
-                                           (or end (point-max)))
-           'type 'variable)))
+  (or (zsh-ts-mode--extract-variable-from-face-at)
      (when-let* ((nodes (treesit-query-capture (treesit-buffer-root-node)
                                                zsh-ts-mode--identifier-query
                                                (point) (1+ (point)))))
-        (if-let* ((var (alist-get 'var nodes)))
-            (propertize (treesit-node-text var t) 'type 'variable)
-          (when-let* ((func (alist-get 'func nodes)))
-            (propertize (treesit-node-text func t) 'type 'function))))))
+        (if-let* ((func (alist-get 'func nodes)))
+            (propertize (treesit-node-text func t) 'type 'function)
+          (when-let* ((var (alist-get 'var nodes)))
+            (propertize (treesit-node-text var t) 'type 'variable))))))

 (cl-defmethod xref-backend-identifier-completion-table ((_backend (eql zsh-ts)))
  "Return nil."
@@ -353,6 +377,8 @@ Return non-nil if NODE is not an anonymous function."
  (treesit-query-compile
   'zsh `((variable_assignment name: (variable_name) @non-let)
          (declaration_command argument: (word) @non-let)
+          (binary_expression left: (variable_name) @non-let
+                             operator: "=")
          (command name: ((command_name) @let.name
                          (:equal @let.name "let"))
                   argument: ((word) @let-arg
@@ -405,11 +431,24 @@ Return non-nil if NODE is not an anonymous function."
             when (save-excursion
                    (goto-char (treesit-node-start node))
                    (looking-at-p quote-ident))
-             collect (zsh-ts-mode--make-xref-for-node node))))
+             collect (zsh-ts-mode--make-xref-for-node node ))))
+
+(defvar zsh-ts-mode--string-query
+  (treesit-query-compile 'zsh '((string) @s))
+  "Query for `zsh-ts-mode--node-not-in-string-p'.")
+
+(defun zsh-ts-mode--node-in-string-p (node)
+  "Return non-nil if NODE is in a string."
+  (treesit-query-capture (treesit-buffer-root-node)
+                         zsh-ts-mode--string-query
+                         (treesit-node-start node)
+                         (treesit-node-start node)))

 (defvar zsh-ts-mode--variable-references-query
  (treesit-query-compile
-   'zsh '(([(simple_variable_name) (special_variable_name)] @var)))
+   'zsh '(([(simple_variable_name) (special_variable_name)] @var)
+          ((arithmetic_expansion :anchor "$((") @arith)
+          ((arithmetic_expansion :anchor "((") @arith-check)))
  "Query used by `zsh-ts-mode--xref-variable-references'.")

 (defun zsh-ts-mode--xref-variable-references (identifier)
@@ -418,14 +457,21 @@ Return non-nil if NODE is not an anonymous function."
        (quote-ident (regexp-quote identifier)))
    (set-text-properties 0 (length ident-copy) () ident-copy)
    (mapcan
-     #'(lambda (node)
-         (when (save-excursion
-                 (goto-char (treesit-node-start node))
-                 (looking-at-p quote-ident))
-           (list (zsh-ts-mode--make-xref-for-node node))))
+     #'(lambda (ent)
+         (cl-destructuring-bind (tag . node) ent
+           (cl-case tag
+            (var (when (save-excursion
+                         (goto-char (treesit-node-start node))
+                         (looking-at-p quote-ident))
+                   (list (zsh-ts-mode--make-xref-for-node node))))
+            (arith (mapcar #'zsh-ts-mode--make-xref-for-node
+                           (zsh-ts-mode--arithmatic-variables node nil)))
+            (arith-check
+             (unless (zsh-ts-mode--node-in-string-p node)
+               (mapcar #'zsh-ts-mode--make-xref-for-node
+                       (zsh-ts-mode--arithmatic-variables node nil)))))))
     (treesit-query-capture (treesit-buffer-root-node)
-                            zsh-ts-mode--variable-references-query
-                            nil nil t))))
+                            zsh-ts-mode--variable-references-query))))

 (cl-defmethod xref-backend-references ((_backend (eql zsh-ts)) identifier)
  "`zsh-ts-mode' implementation for finding xref references for IDENTIFIER."
@@ -443,11 +489,14 @@ Return non-nil if NODE is not an anonymous function."
   'zsh '(((simple_variable_name) @ident)
          ((special_variable_name) @ident)
          ((variable_name) @ident)
-          (command name: (command_name ( (word) @cname (:equal @cname "let")))
+          (command name: (command_name ((word) @cname (:equal @cname "let")))
                   argument: (word) @ident)
          (declaration_command argument: (word) @ident)
+          (command name: (command_name (word) @ident))
          ((variable_name) @ident)
-          (function_definition name: (word) @ident)))
+          (function_definition name: (word) @ident)
+          ((arithmetic_expansion :anchor "$((") @arith)
+          ((arithmetic_expansion :anchor "((") @arith-check)))
  "Query used by `xref-backend-apropos'.")

 (cl-defmethod xref-backend-apropos ((_backend (eql zsh-ts)) pattern)
@@ -456,9 +505,14 @@ Return non-nil if NODE is not an anonymous function."
    (dolist (ent (treesit-query-capture (treesit-buffer-root-node)
                                        zsh-ts-mode--xref-apropos-query))
      (cl-destructuring-bind (type . node) ent
-        (when (and (eq type 'ident)
-                   (string-match-p pattern (treesit-node-text node)))
-          (push (zsh-ts-mode--make-xref-for-node node) out))))
+        (cond
+         ((eq type 'ident) (when (string-match-p pattern (treesit-node-text node))
+                             (push (zsh-ts-mode--make-xref-for-node node) out)))
+         ((or (eq type 'arith)
+              (not (zsh-ts-mode--node-in-string-p node)))
+          (dolist (var-node (zsh-ts-mode--arithmatic-variables node nil))
+            (when (string-match-p pattern (treesit-node-text var-node))
+              (push (zsh-ts-mode--make-xref-for-node var-node) out)))))))
    out))

 ;;;###autoload
@@ -488,5 +542,19 @@ not written in Zsh."
    (add-hook 'xref-backend-functions #'zsh-ts-mode--xref-backend nil t)
    (treesit-major-mode-setup)))

+(derived-mode-add-parents 'zsh-ts-mode '(sh-mode))
+
+;;;###autoload
+(add-to-list 'auto-mode-alist
+             `(,(rx (or ".zsh"
+                        (seq (or "/" bos)
+                             (or ".zshrc" ".zprofile" ".zlogin"
+                                 ".zshenv")))
+                    eos)
+               . zsh-ts-mode))
+
+;;;###autoload
+(add-to-list 'interpreter-mode-alist '("zsh" . zsh-ts-mode))
+
 (provide 'zsh-ts-mode)
 ;;; zsh-ts-mode.el ends here
--- a/init.el
+++ b/init.el
@@ -2134,7 +2134,7 @@ This is :around advice, so OLDFUN is the real function
  :config
  (my/setup-c-style-newline-keys typescript-ts-mode-map))

-;; shell-mode
+;; sh-mode
 (use-package sh-script
  :ensure nil
  :hook (sh-mode . my/-setup-sh-mode)
@@ -2143,6 +2143,10 @@ This is :around advice, so OLDFUN is the real function
  (defun my/-setup-sh-mode ()
    (add-hook 'completion-at-point-functions #'cape-file nil t)))

+;; zsh-ts-mode
+(use-package zsh-ts-mode
+  :ensure nil)
+
 ;; go mode
 (use-package go-mode
  :defer nil