From 795e01ac248d389a581589b13a02465a2f99202f Mon Sep 17 00:00:00 2001 From: Yuan Fu Date: Sat, 24 Sep 2022 19:41:17 -0700 Subject: [PATCH] Update and enable treesit-imenu function in python.el * lisp/progmodes/python.el (python--treesit-settings): Add docstring. (python--imenu-treesit-create-index-1): Rewrite with treesit-induce-sparse-tree. (python-imenu-treesit-create-index): Move main body to python--imenu-treesit-create-index-1. (python-imenu-treesit-create-flat-index): Fix typo. (python-mode): Enable treesit-imenu. Also fix indentation for which-func code. --- lisp/progmodes/python.el | 114 +++++++++++++++++++++++---------------- 1 file changed, 69 insertions(+), 45 deletions(-) diff --git a/lisp/progmodes/python.el b/lisp/progmodes/python.el index 8368f4da513..fb91d00053a 100644 --- a/lisp/progmodes/python.el +++ b/lisp/progmodes/python.el @@ -1035,7 +1035,9 @@ Do not fontify the initial f for f-strings." `(seq bol (or ,@python--treesit-exceptions) eol)) @font-lock-type-face)) - (type (identifier) @font-lock-type-face)))) + (type (identifier) @font-lock-type-face))) + "Tree-sitter font-lock settings.") + ;;; Indentation @@ -5244,61 +5246,79 @@ To this: (python-imenu-create-index)))))) ;;; Tree-sitter imenu -;; -;; This works, but is slower than the native functions, presumably -;; because traversing the parser tree is slower than scanning the -;; text. Also I'm sure this consumes more memory as we allocate -;; memory for every node in the tree. -(defun python--imenu-treesit-create-index (&optional node) - "Return tree Imenu alist for the current Python buffer. +(defun python--imenu-treesit-create-index-1 (node) + "Given a sparse tree, create an imenu alist. -Change `python-imenu-format-item-label-function', -`python-imenu-format-parent-item-label-function', -`python-imenu-format-parent-item-jump-label-function' to -customize how labels are formatted. +NODE is the root node of the tree returned by +`treesit-induce-sparse-tree' (not a tree-sitter node, its car is +a tree-sitter node). Walk that tree and return an imenu alist. -NODE is the root node of the subtree you want to build an index -of. If nil, use the root node of the whole parse tree. +Return a list of ENTRY where -Similar to `python-imenu-create-index' but use tree-sitter." - (let* ((node (or node (treesit-buffer-root-node 'python))) - (children (treesit-node-children node t)) - (subtrees (mapcan #'python--imenu-treesit-create-index +ENTRY := (NAME . MARKER) + | (NAME . ((JUMP-LABEL . MARKER) + ENTRY + ...) + +NAME is the function/class's name, JUMP-LABEL is like \"*function +definition*\"." + (let* ((ts-node (car node)) + (children (cdr node)) + (subtrees (mapcan #'python--imenu-treesit-create-index-1 children)) - (type (pcase (treesit-node-type node) - ("function_definition" 'def) - ("class_definition" 'class) - (_ nil))) - (name (when type + (type (pcase (treesit-node-type ts-node) + ("function_definition" 'def) + ("class_definition" 'class))) + ;; The root of the tree could have a nil ts-node. + (name (when ts-node (treesit-node-text (treesit-node-child-by-field-name - node "name") t)))) + ts-node "name") t))) + (marker (when ts-node + (set-marker (make-marker) + (treesit-node-start ts-node))))) (cond - ;; 1. This node is a function/class and doesn't have children. - ((and type (not subtrees)) - (let ((label - (funcall python-imenu-format-item-label-function - type name))) - (list (cons label - (set-marker (make-marker) - (treesit-node-start node)))))) - ;; 2. This node is a function/class and has children. - ((and type subtrees) + ((null ts-node) + subtrees) + (subtrees (let ((parent-label (funcall python-imenu-format-parent-item-label-function type name)) (jump-label - (funcall python-imenu-format-parent-item-jump-label-function - type name))) + (funcall + python-imenu-format-parent-item-jump-label-function + type name))) `((,parent-label - ,(cons jump-label (set-marker (make-marker) - (treesit-node-start node))) + ,(cons jump-label marker) ,@subtrees)))) - ;; 3. This node is not a function/class. - ((not type) subtrees)))) + (t (let ((label + (funcall python-imenu-format-item-label-function + type name))) + (list (cons label marker))))))) + +(defun python-imenu-treesit-create-index (&optional node) + "Return tree Imenu alist for the current Python buffer. -(defun python--imenu-treesit-create-flat-index () +Change `python-imenu-format-item-label-function', +`python-imenu-format-parent-item-label-function', +`python-imenu-format-parent-item-jump-label-function' to +customize how labels are formatted. + +NODE is the root node of the subtree you want to build an index +of. If nil, use the root node of the whole parse tree. + +Similar to `python-imenu-create-index' but use tree-sitter." + (let* ((node (or node (treesit-buffer-root-node 'python))) + (tree (treesit-induce-sparse-tree + node + (rx (seq bol + (or "function" "class") + "_definition" + eol))))) + (python--imenu-treesit-create-index-1 tree))) + +(defun python-imenu-treesit-create-flat-index () "Return flat outline of the current Python buffer for Imenu. Change `python-imenu-format-item-label-function', @@ -5309,7 +5329,7 @@ customize how labels are formatted. Similar to `python-imenu-create-flat-index' but use tree-sitter." (python-imenu-create-flat-index - (python--imenu-treesit-create-index))) + (python-imenu-treesit-create-index))) ;;; Misc helpers @@ -6120,14 +6140,18 @@ REPORT-FN is Flymake's callback function." (add-hook 'post-self-insert-hook #'python-indent-post-self-insert-function 'append 'local) - (setq-local imenu-create-index-function - #'python-imenu-create-index) + (if (and python-use-tree-sitter + (treesit-can-enable-p)) + (setq-local imenu-create-index-function + #'python-treesit-imenu-create-index) + (setq-local imenu-create-index-function + #'python-imenu-create-index)) (setq-local add-log-current-defun-function #'python-info-current-defun) (if (and python-use-tree-sitter - (treesit-can-enable-p)) + (treesit-can-enable-p)) (add-hook 'which-func-functions #'python-info-treesit-current-defun nil t) (add-hook 'which-func-functions #'python-info-current-defun nil t)) -- 2.39.5