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.
This commit is contained in:
Yuan Fu 2022-11-08 20:41:58 -08:00
parent 5858921f40
commit 2b4d46f99b
No known key found for this signature in database
GPG key ID: 56E19BC57664A442

View file

@ -2353,6 +2353,55 @@ 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
@ -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