From 936831579490c2e2a057298f5f915465fbb116d8 Mon Sep 17 00:00:00 2001
From: Yuan Fu <casouri@gmail.com>
Date: Wed, 23 Nov 2022 12:08:47 -0800
Subject: [PATCH] Don't skip nested defuns in python-ts-mode defun navigation

This fixes bug#59495.

Before this change, python tries to skip nested function
definition. Now we don't skip any nested defun.

* lisp/progmodes/python.el (python-treesit-beginning-of-defun)
(python-treesit-end-of-defun): Remove functions.
* lisp/progmodes/python.el (python-ts-mode): Use tree-sitter's default
navigation function.
---
 lisp/progmodes/python.el | 54 ++--------------------------------------
 1 file changed, 2 insertions(+), 52 deletions(-)

diff --git a/lisp/progmodes/python.el b/lisp/progmodes/python.el
index 18594a3e23d..f97ae81508a 100644
--- a/lisp/progmodes/python.el
+++ b/lisp/progmodes/python.el
@@ -2395,55 +2395,6 @@ position, else returns nil."
         (point)
       (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
 
@@ -6655,9 +6606,8 @@ implementations: `python-mode' and `python-ts-mode'."
     (setq-local treesit-font-lock-settings python--treesit-settings)
     (setq-local imenu-create-index-function
                 #'python-imenu-treesit-create-index)
-    (setq-local beginning-of-defun-function
-                #'python-treesit-beginning-of-defun)
-    (setq-local end-of-defun-function #'python-treesit-end-of-defun)
+    (setq-local treesit-defun-type-regexp (rx (or "function" "class")
+                                              "_definition"))
     (treesit-major-mode-setup)))
 
 ;;; Completion predicates for M-x
-- 
2.39.5