Use tree sitter for getting the node id (aka test id) for the function and class at point

python-pytest.el

@@ -1,7 +1,7 @@
 ;;; python-pytest.el --- helpers to run pytest -*- lexical-binding: t; -*-
 
 ;; Author: wouter bolsterlee <[email protected]>
-;; Version: 3.3.0
+;; Version: 3.5.0
 ;; Package-Requires: ((emacs "24.4") (dash "2.18.0") (transient "0.3.7") (s "1.12.0"))
 ;; Keywords: pytest, test, python, languages, processes, tools
 ;; URL: https://github.com/wbolster/emacs-python-pytest
@@ -29,6 +29,7 @@
 
 (require 'projectile nil t)
 (require 'project nil t)
+(require 'treesit nil t)
 
 (defgroup python-pytest nil
   "pytest integration"
@@ -127,6 +128,16 @@ When non-nil only ‘test_foo()’ will match, and nothing else."
            (set-default symbol value)
            value))))
 
+(defcustom python-pytest-use-treesit (featurep 'treesit)
+  "Whether to use treesit for getting the node ids of things at point.
+
+Users that are running a version of Emacs that supports treesit
+and have the Python language grammar for treesit should set this
+variable to t. Users that are running a version of Emacs that
+don't support treesit should set this variable to nil."
+  :group 'python-pytest
+  :type 'boolean)
+
 (defvar python-pytest--history nil
   "History for pytest invocations.")
 
@@ -178,8 +189,10 @@ When non-nil only ‘test_foo()’ will match, and nothing else."
     ("F" "file (this)" python-pytest-file)]
    [("m" "files" python-pytest-files)
     ("M" "directories" python-pytest-directories)]
-   [("d" "def/class (dwim)" python-pytest-function-dwim)
-    ("D" "def/class (this)" python-pytest-function)]])
+   [("d" "def at point (dwim)" python-pytest-run-def-or-class-at-point-dwim :if-not python-pytest--use-treesit-p)
+    ("D" "def at point" python-pytest-run-def-or-class-at-point :if-not python-pytest--use-treesit-p)
+    ("d" "def at point" python-pytest-run-def-at-point-treesit :if python-pytest--use-treesit-p)
+    ("c" "class at point" python-pytest-run-class-at-point-treesit :if python-pytest--use-treesit-p)]])
 
 (define-obsolete-function-alias 'python-pytest-popup 'python-pytest-dispatch "2.0.0")
 
@@ -258,35 +271,60 @@ With a prefix argument, allow editing."
    :edit current-prefix-arg))
 
 ;;;###autoload
-(defun python-pytest-function (file func args)
+(defun python-pytest-run-def-at-point-treesit ()
+  "Run def at point."
+  (interactive)
+  (python-pytest--run
+   :args (transient-args 'python-pytest-dispatch)
+   :file (buffer-file-name)
+   :node-id (python-pytest--node-id-def-at-point-treesit)
+   :edit current-prefix-arg))
+
+;;;###autoload
+(defun python-pytest-run-class-at-point-treesit ()
+  "Run class at point."
+  (interactive)
+  (python-pytest--run
+   :args (transient-args 'python-pytest-dispatch)
+   :file (buffer-file-name)
+   :node-id (python-pytest--node-id-class-at-point-treesit)
+   :edit current-prefix-arg))
+
+;;;###autoload
+(defun python-pytest-run-def-or-class-at-point (file func args)
   "Run pytest on FILE with FUNC (or class).
 
 Additional ARGS are passed along to pytest.
 With a prefix argument, allow editing."
   (interactive
    (list
     (buffer-file-name)
-    (python-pytest--current-defun)
+    (python-pytest--node-id-def-or-class-at-point)
     (transient-args 'python-pytest-dispatch)))
   (python-pytest--run
    :args args
    :file file
-   :func func
+   :node-id func
    :edit current-prefix-arg))
 
 ;;;###autoload
-(defun python-pytest-function-dwim (file func args)
-  "Run pytest on FILE with FUNC (or class).
+(defun python-pytest-run-def-or-class-at-point-dwim (file func args)
+  "Run pytest on FILE using FUNC at point as the node-id.
 
-When run interactively, this tries to work sensibly using
-the current file and function around point.
+If `python-pytest--test-file-p' returns t for FILE (i.e. the file
+is a test file), then this function results in the same behavior
+as calling `python-pytest-run-def-at-point'. If
+`python-pytest--test-file-p' returns nil for FILE (i.e. the
+current file is not a test file), then this function will try to
+find related test files and test defs (i.e. sensible match) for
+the current file and the def at point.
 
 Additional ARGS are passed along to pytest.
 With a prefix argument, allow editing."
   (interactive
    (list
     (buffer-file-name)
-    (python-pytest--current-defun)
+    (python-pytest--node-id-def-or-class-at-point)
     (transient-args 'python-pytest-dispatch)))
   (unless (python-pytest--test-file-p file)
     (setq
@@ -313,7 +351,7 @@ With a prefix argument, allow editing."
   (python-pytest--run
    :args args
    :file file
-   :func func
+   :node-id func
    :edit current-prefix-arg))
 
 ;;;###autoload
@@ -360,16 +398,22 @@ With a prefix ARG, allow editing."
     map)
   "Keymap for `python-pytest-mode' major mode.")
 
-(cl-defun python-pytest--run (&key args file func edit)
-  "Run pytest for the given arguments."
+(cl-defun python-pytest--run (&key args file node-id edit)
+  "Run pytest for the given arguments.
+
+NODE-ID should be the node id of the test to run. pytest uses
+double colon \"::\" for separating components in node ids. For
+example, the node-id for a function outside a class is the
+function name, the node-id for a function inside a class is
+TestClass::test_my_function, the node-id for a function inside a
+class that is inside another class is
+TestClassParent::TestClassChild::test_my_function."
   (setq args (python-pytest--transform-arguments args))
   (when (and file (file-name-absolute-p file))
     (setq file (python-pytest--relative-file-name file)))
-  (when func
-    (setq func (s-replace "." "::" func)))
   (let ((command)
         (thing (cond
-                ((and file func) (format "%s::%s" file func))
+                ((and file node-id) (format "%s::%s" file node-id))
                 (file file))))
     (when thing
       (setq args (-snoc args (python-pytest--shell-quote thing))))
@@ -429,6 +473,17 @@ With a prefix ARG, allow editing."
       (setq process (get-buffer-process buffer))
       (set-process-sentinel process #'python-pytest--process-sentinel))))
 
+(defun python-pytest--use-treesit-p ()
+  "Return t if python-pytest-use-treesit is t. Otherwise, return nil.
+
+This function is passed to the parameter :if in
+`python-pytest-dispatch'.
+
+Although this function might look useless, the main reason why it
+was defined was that the parameter that is provided to the
+transient keyword :if must be a function."
+  python-pytest-use-treesit)
+
 (defun python-pytest--shell-quote (s)
   "Quote S for use in a shell command. Like `shell-quote-argument', but prettier."
   (if (s-equals-p s (shell-quote-argument s))
@@ -545,7 +600,140 @@ When present ON-REPLACEMENT is substituted, else OFF-REPLACEMENT is appended."
 
 ;; python helpers
 
-(defun python-pytest--current-defun ()
+(defun python-pytest--point-is-inside-def-treesit ()
+  (unless (treesit-language-available-p 'python)
+    (error "This function requires tree-sitter support for python, but it is not available."))
+  (save-restriction
+    (widen)
+    (catch 'return
+      (let ((current-node (treesit-node-at (point) 'python)))
+        (while (setq current-node (treesit-node-parent current-node))
+          (when (equal (treesit-node-type current-node) "function_definition")
+            (throw 'return t)))))))
+
+(defun python-pytest--point-is-inside-class-treesit ()
+  (unless (treesit-language-available-p 'python)
+    (error "This function requires tree-sitter support for python, but it is not available."))
+  (save-restriction
+    (widen)
+    (catch 'return
+      (let ((current-node (treesit-node-at (point) 'python)))
+        (while (setq current-node (treesit-node-parent current-node))
+          (when (equal (treesit-node-type current-node) "class_definition")
+            (throw 'return t)))))))
+
+(defun python-pytest--node-id-def-at-point-treesit ()
+  "Return the node id of the def at point.
+
++ If the test function is not inside a class, its node id is the name
+  of the function.
++ If the test function is defined inside a class, its node id would
+  look like: TestGroup::test_my_function.
++ If the test function is defined inside a class that is defined
+  inside another class, its node id would look like:
+  TestGroupParent::TestGroupChild::test_my_function."
+  (unless (python-pytest--point-is-inside-def-treesit)
+    (error "The point is not inside a def."))
+  (save-restriction
+    (widen)
+    (let ((function
+           ;; Move up to the outermost function
+           (catch 'return
+             (let ((current-node (treesit-node-at (point) 'python))
+                   function-node)
+               (catch 'break
+                 (while (setq current-node (treesit-node-parent current-node))
+                   (when (equal (treesit-node-type current-node) "function_definition")
+                     (setq function-node current-node)
+                     ;; At this point, we know that we are on a
+                     ;; function. We need to move up to see if the
+                     ;; function is inside a function. If that's the
+                     ;; case, we move up. This way, we find the
+                     ;; outermost function. We need to do this because
+                     ;; pytest can't execute functions inside functions,
+                     ;; so we must get the function that is not inside
+                     ;; other function.
+                     (while (setq current-node (treesit-node-parent current-node))
+                       (when (equal (treesit-node-type current-node) "function_definition")
+                         (setq function-node current-node)))
+                     (throw 'break nil))))
+               (dolist (child (treesit-node-children function-node))
+                 (when (equal (treesit-node-type child) "identifier")
+                   (throw 'return
+                          (cons
+                           ;; Keep a reference to the node that is a
+                           ;; function_definition. We need this
+                           ;; reference because afterwards we need to
+                           ;; move up starting at the current node to
+                           ;; find the node id of the class (if there's
+                           ;; any) in which the function is defined.
+                           function-node
+                           (buffer-substring-no-properties
+                            (treesit-node-start child)
+                            (treesit-node-end child)))))))))
+          parents)
+      ;; Move up through the parent nodes to see if the function is
+      ;; defined inside a class and collect the classes to finally build
+      ;; the node id of the current function. Remember that the node id
+      ;; of a function that is defined within nested classes must have
+      ;; the name of the nested classes.
+      (let ((current-node (car function)))
+        (while (setq current-node (treesit-node-parent current-node))
+          (when (equal (treesit-node-type current-node) "class_definition")
+            (dolist (child (treesit-node-children current-node))
+              (when (equal (treesit-node-type child) "identifier")
+                (push (buffer-substring-no-properties
+                       (treesit-node-start child)
+                       (treesit-node-end child))
+                      parents))))))
+      (string-join `(,@parents ,(cdr function)) "::"))))
+
+(defun python-pytest--node-id-class-at-point-treesit ()
+  "Return the node id of the class at point.
+
++ If the class is not inside another class, its node id is the name
+  of the class.
++ If the class is defined inside another class, the node id of the
+  class which is contained would be: TestGroupParent::TestGroupChild,
+  while the node id of the class which contains the other class would
+  be TestGroupParent."
+  (unless (python-pytest--point-is-inside-class-treesit)
+    (error "The point is not inside a class."))
+  (save-restriction
+    (widen)
+    (let ((class
+           ;; Move up to the outermost function
+           (catch 'return
+             (let ((current-node (treesit-node-at (point) 'python)))
+               (catch 'break
+                 (while (setq current-node (treesit-node-parent current-node))
+                   (when (equal (treesit-node-type current-node) "class_definition")
+                     (throw 'break nil))))
+               (dolist (child (treesit-node-children current-node))
+                 (when (equal (treesit-node-type child) "identifier")
+                   (throw 'return
+                          (cons
+                           ;; Keep a reference to the node that is a
+                           ;; function_definition
+                           current-node
+                           (buffer-substring-no-properties
+                            (treesit-node-start child)
+                            (treesit-node-end child)))))))))
+          parents)
+      ;; Move up through the parents to collect the list of classes in
+      ;; which the class is contained. pytest supports running nested
+      ;; classes, but it doesn't support runing nested functions.
+      (let ((current-node (car class)))
+        (while (setq current-node (treesit-node-parent current-node))
+          (when (equal (treesit-node-type current-node) "class_definition")
+            (dolist (child (treesit-node-children current-node))
+              (when (equal (treesit-node-type child) "identifier")
+                (push (buffer-substring-no-properties
+                       (treesit-node-start child)
+                       (treesit-node-end child))
+                      parents))))))
+      (string-join `(,@parents ,(cdr class)) "::"))))
+(defun python-pytest--node-id-def-or-class-at-point ()
   "Detect the current function/class (if any)."
   (let* ((name
           (or (python-info-current-defun)
@@ -565,7 +753,7 @@ When present ON-REPLACEMENT is substituted, else OFF-REPLACEMENT is appended."
           (if (s-lowercase? (substring name 0 1))
               (car (s-split-up-to "\\." name 1))
             name)))
-    name))
+    (s-replace "." "::" name)))
 
 (defun python-pytest--make-test-name (func)
   "Turn function name FUNC into a name (hopefully) matching its test name.

tests/README.org

@@ -0,0 +1,9 @@
+The following command can be used to run all tests in the directory =tests=. The command should be run in the root directory of the project. The command explicitly loads the file =python-pytest.el= in this repository, this is done to make sure that Emacs uses the symbol definitions from that file instead of other locations that might have the same package (e.g. installed through MELPA.)
+
+#+BEGIN_SRC sh
+emacs \
+  --batch \
+  --eval '(load-file "./python-pytest.el")' \
+  --eval '(dolist (file (directory-files-recursively "tests" "\\`[^.].*\\.el\\'\''")) (load-file file))' \
+  --eval '(ert-run-tests-batch-and-exit)'
+#+END_SRC