feat(parsers): Improve folding of block nodes in python (#32)

* feat(parsers): Improve folding of block nodes in python

Add a new method treesit-fold-range-python-block in treesit-fold-parsers-python
to fold below code blocks in python-ts-mode.

- while_statement
- for_statement
- if_statement
- elif_clause
- else_clause
- match_statement
- case_clause
- try_statement
- except_clause
- with_statement

Fixes #31

* style: clean up

* docs: changelog

---------

Co-authored-by: JenChieh <jcs090218@gmail.com>

Author: lijoantony

CHANGELOG.md

@@ -8,7 +8,7 @@ Check [Keep a Changelog](http://keepachangelog.com/) for recommendations on how
 ## 0.3.0 (Unreleased)
 > Released N/A
 
-* N/A
+* feat(parsers): Improve folding of block nodes in python (#32)
 
 ## 0.2.0
 > Released Feb 12, 2025

treesit-fold-parsers.el

@@ -90,6 +90,7 @@
 (declare-function treesit-fold-range-cmake-body "treesit-fold.el")
 (declare-function treesit-fold-range-editorconfig-section "treesit-fold.el")
 (declare-function treesit-fold-range-pascal-comment "treesit-fold.el")
+(declare-function treesit-fold-range-python-block "treesit-fold.el")
 (declare-function treesit-fold-range-python-def "treesit-fold.el")
 (declare-function treesit-fold-range-python-expression-statement "treesit-fold.el")
 (declare-function treesit-fold-range-rst-body "treesit-fold.el")
@@ -587,6 +588,16 @@
     (future_import_statement  . treesit-fold-range-seq)
     (function_definition      . treesit-fold-range-python-def)
     (class_definition         . treesit-fold-range-python-def)
+    (while_statement          . treesit-fold-range-python-block)
+    (for_statement            . treesit-fold-range-python-block)
+    (if_statement             . treesit-fold-range-python-block)
+    (elif_clause              . treesit-fold-range-python-block)
+    (else_clause              . treesit-fold-range-python-block)
+    (match_statement          . treesit-fold-range-python-block)
+    (case_clause              . treesit-fold-range-python-block)
+    (try_statement            . treesit-fold-range-python-block)
+    (except_clause            . treesit-fold-range-python-block)
+    (with_statement           . treesit-fold-range-python-block)
     (list                     . treesit-fold-range-seq)
     (dictionary               . treesit-fold-range-seq)
     (parenthesized_expression . treesit-fold-range-seq)

treesit-fold.el

@@ -1399,6 +1399,34 @@ more information."
           (t
            (treesit-fold-range-c-like-comment node offset)))))
 
+(defun treesit-fold-range-python-block (node offset)
+  "Define fold range for `if_statement' and other blocks.
+
+For arguments NODE and OFFSET, see function `treesit-fold-range-seq' for
+more information."
+  (when-let* ((colon-node (car (treesit-fold-find-children node ":")))
+              (beg (treesit-node-start colon-node)))
+    (let ((current-node (treesit-node-next-sibling colon-node))
+          (last-body-node)
+          (end))
+      ;; Iterate through siblings until we hit an elif or else clause
+      (while (and current-node
+                  (not (member (treesit-node-type current-node)
+                               '("elif_clause" "else_clause" "except_clause"))))
+
+        ;; Only consider non-comment nodes as body nodes
+        (unless (string-match-p "comment" (treesit-node-type current-node))
+          (setq last-body-node current-node))
+        (setq current-node (treesit-node-next-sibling current-node)))
+
+      ;; Set end position based on the last body node or fallback to node end
+      (setq end (if last-body-node
+                    (treesit-node-end last-body-node)
+                  (treesit-node-end node)))
+
+      ;; Return the range and offset to fold
+      (treesit-fold--cons-add (cons (+ beg 1) end) offset))))
+
 (defun treesit-fold-range-python-def (node offset)
   "Define fold range for `function_definition' and `class_definition'.