Mimic existing python-mode beg/end-of-defun behavior better

* lisp/progmodes/python.el (python-treesit-beginning-of-defun)
(python-treesit-end-of-defun): New functions.
* lisp/progmodes/python.el (python-mode): Use custom beg/end-of-defun
functions.

1 files changed, 52 insertions, 4 deletions

diff --git a/lisp/progmodes/python.el b/lisp/progmodes/python.el
index 8db96b117f3..61f29dd005c 100644
--- a/lisp/progmodes/python.el
+++ b/lisp/progmodes/python.el
@@ -2354,6 +2354,55 @@ position, else returns nil."
       (ignore (goto-char point)))))
 
 
+;;; Tree-sitter navigation
+
+(defun python-treesit-beginning-of-defun (&optional arg)
+  "Tree-sitter `beginning-of-defun' function.
+ARG is the same as in `beginning-of-defun'."
+  (let ((arg (or arg 1))
+        (node (treesit-node-at (point)))
+        (function-or-class (rx (or "function" "class") "_definition")))
+    (if (> arg 0)
+        ;; Go backward.
+        (while (and (> arg 0)
+                    (setq node (treesit-search-forward-goto
+                                node function-or-class t t)))
+          ;; Here we deviate from `treesit-beginning-of-defun': if
+          ;; NODE is function_definition, find the top-level
+          ;; function_definition, if NODE is class_definition, find
+          ;; the top-level class_definition, don't mix the two like
+          ;; `treesit-beginning-of-defun' would.
+          (setq node (or (treesit-node-top-level node)
+                         node))
+          (setq arg (1- arg)))
+      ;; Go forward.
+      (while (and (< arg 0)
+                  (setq node (treesit-search-forward-goto
+                              node function-or-class)))
+        (setq node (or (treesit-node-top-level node)
+                       node))
+        (setq arg (1+ arg))))
+    (when node
+      (goto-char (treesit-node-start node))
+      t)))
+
+(defun python-treesit-end-of-defun ()
+  "Tree-sitter `end-of-defun' function."
+  ;; Why not simply get the largest node at point: when point is at
+  ;; (point-min), that gives us the root node.
+  (let* ((node (treesit-node-at (point)))
+         (top-func (treesit-node-top-level
+                    node
+                    "function_definition"))
+         (top-class (treesit-node-top-level
+                     node
+                     "class_definition")))
+    ;; Prefer function_definition over class_definition: when we are
+    ;; in a function_definition inside a class_definition, jump to the
+    ;; end of function_definition.
+    (goto-char (or (treesit-node-end (or top-func top-class)) (point)))))
+
+
 ;;; Shell integration
 
 (defcustom python-shell-buffer-name "Python"
@@ -6508,10 +6557,9 @@ Add import for undefined name `%s' (empty to skip): "
     (setq-local treesit-font-lock-settings python--treesit-settings)
     (setq-local imenu-create-index-function
                 #'python-imenu-treesit-create-index)
-    (setq-local treesit-defun-type-regexp (rx bol
-                                              (or "function" "class")
-                                              "_definition"
-                                              eol))
+    (setq-local beginning-of-defun-function
+                #'python-treesit-beginning-of-defun)
+    (setq-local end-of-defun-function #'python-treesit-end-of-defun)
     (treesit-major-mode-setup))
    ;; Elisp.
    (t