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.

1 files 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))