2 files changed, 247 insertions, 7 deletions

diff --git a/lisp/progmodes/python.el b/lisp/progmodes/python.el
index e1347754c4a..d3ffc2db2c9 100644
--- a/lisp/progmodes/python.el
+++ b/lisp/progmodes/python.el
@@ -1524,6 +1524,10 @@ marks the next defun after the ones already marked."
 The name of the defun should be grouped so it can be retrieved
 via `match-string'.")
 
+(defvar python-nav-beginning-of-block-regexp
+  (python-rx line-start (* space) block-start)
+  "Regexp matching block start.")
+
 (defun python-nav--beginning-of-defun (&optional arg)
   "Internal implementation of `python-nav-beginning-of-defun'.
 With positive ARG search backwards, else search forwards."
@@ -4916,9 +4920,37 @@ Interactively, prompt for symbol."
 (defun python-hideshow-forward-sexp-function (_arg)
   "Python specific `forward-sexp' function for `hs-minor-mode'.
 Argument ARG is ignored."
-  (python-nav-end-of-defun)
-  (unless (python-info-current-line-empty-p)
-    (backward-char)))
+  (python-nav-end-of-block))
+
+(defun python-hideshow-find-next-block (regexp maxp comments)
+  "Python specific `hs-find-next-block' function for `hs-minor-mode'.
+Call `python-nav-forward-block' to find next block and check if
+block-start ends within MAXP.  If COMMENTS is not nil, comments
+are also searched.  REGEXP is passed to `looking-at' to set
+`match-data'."
+  (let* ((next-block (save-excursion
+                       (or (and
+                            (python-info-looking-at-beginning-of-block)
+                            (re-search-forward
+                             (python-rx block-start) maxp t))
+                           (and (python-nav-forward-block)
+                                (< (point) maxp)
+                                (re-search-forward
+                                 (python-rx block-start) maxp t))
+                           (1+ maxp))))
+         (next-comment
+          (or (when comments
+                (save-excursion
+                  (cl-loop while (re-search-forward "#" maxp t)
+                           if (python-syntax-context 'comment)
+                           return (point))))
+              (1+ maxp)))
+         (next-block-or-comment (min next-block next-comment)))
+    (when (<= next-block-or-comment maxp)
+      (goto-char next-block-or-comment)
+      (save-excursion
+        (beginning-of-line)
+        (looking-at regexp)))))
 
 
 ;;; Imenu
@@ -5415,6 +5447,16 @@ instead of the current physical line."
          (beginning-of-line 1)
          (looking-at python-nav-beginning-of-defun-regexp))))
 
+(defun python-info-looking-at-beginning-of-block ()
+  "Check if point is at the beginning of block."
+  (let ((pos (point)))
+    (save-excursion
+      (python-nav-beginning-of-statement)
+      (beginning-of-line)
+      (and
+       (<= (point) pos (+ (point) (current-indentation)))
+       (looking-at python-nav-beginning-of-block-regexp)))))
+
 (defun python-info-current-line-comment-p ()
   "Return non-nil if current line is a comment line."
   (char-equal
@@ -5870,14 +5912,17 @@ REPORT-FN is Flymake's callback function."
 
   (add-to-list
    'hs-special-modes-alist
-   '(python-mode
-     "\\s-*\\_<\\(?:def\\|class\\)\\_>"
+   `(python-mode
+     ,python-nav-beginning-of-block-regexp
      ;; Use the empty string as end regexp so it doesn't default to
      ;; "\\s)".  This way parens at end of defun are properly hidden.
      ""
      "#"
      python-hideshow-forward-sexp-function
-     nil))
+     nil
+     python-nav-beginning-of-block
+     python-hideshow-find-next-block
+     python-info-looking-at-beginning-of-block))
 
   (setq-local outline-regexp (python-rx (* space) block-start))
   (setq-local outline-level
diff --git a/test/lisp/progmodes/python-tests.el b/test/lisp/progmodes/python-tests.el
index 12ac871fdf8..906f7eca7de 100644
--- a/test/lisp/progmodes/python-tests.el
+++ b/test/lisp/progmodes/python-tests.el
@@ -5730,6 +5730,39 @@ def \\
    (should (not (python-info-looking-at-beginning-of-defun)))
    (should (not (python-info-looking-at-beginning-of-defun nil t)))))
 
+(ert-deftest python-info-looking-at-beginning-of-block-1 ()
+  (python-tests-with-temp-buffer
+   "
+def f():
+    if True:
+        pass
+    l = [x * 2
+         for x in range(5)
+         if x < 3]
+# if False:
+\"\"\"
+if 0:
+\"\"\"
+"
+   (python-tests-look-at "def f():")
+   (should (python-info-looking-at-beginning-of-block))
+   (forward-char)
+   (should (not (python-info-looking-at-beginning-of-block)))
+   (python-tests-look-at "if True:")
+   (should (python-info-looking-at-beginning-of-block))
+   (forward-char)
+   (should (not (python-info-looking-at-beginning-of-block)))
+   (beginning-of-line)
+   (should (python-info-looking-at-beginning-of-block))
+   (python-tests-look-at "for x")
+   (should (not (python-info-looking-at-beginning-of-block)))
+   (python-tests-look-at "if x < 3")
+   (should (not (python-info-looking-at-beginning-of-block)))
+   (python-tests-look-at "if False:")
+   (should (not (python-info-looking-at-beginning-of-block)))
+   (python-tests-look-at "if 0:")
+   (should (not (python-info-looking-at-beginning-of-block)))))
+
 (ert-deftest python-info-current-line-comment-p-1 ()
   (python-tests-with-temp-buffer
    "
@@ -6183,8 +6216,11 @@ class SomeClass:
 class SomeClass:
 
     def __init__(self, arg, kwarg=1):
+
     def filter(self, nums):
-    def __str__(self):"))))
+
+    def __str__(self):
+"))))
       (or enabled (hs-minor-mode -1)))))
 
 (ert-deftest python-hideshow-hide-levels-2 ()
@@ -6230,6 +6266,165 @@ class SomeClass:
 "))))
       (or enabled (hs-minor-mode -1)))))
 
+(ert-deftest python-hideshow-hide-levels-3 ()
+  "Should hide all blocks."
+  (python-tests-with-temp-buffer
+   "
+def f():
+    if 0:
+        l = [i for i in range(5)
+             if i < 3]
+        abc = o.match(1, 2, 3)
+
+def g():
+    pass
+"
+   (hs-minor-mode 1)
+   (hs-hide-level 1)
+   (should
+    (string=
+     (python-tests-visible-string)
+     "
+def f():
+
+def g():
+"))))
+
+(ert-deftest python-hideshow-hide-levels-4 ()
+  "Should hide 2nd level block."
+  (python-tests-with-temp-buffer
+   "
+def f():
+    if 0:
+        l = [i for i in range(5)
+             if i < 3]
+        abc = o.match(1, 2, 3)
+
+def g():
+    pass
+"
+   (hs-minor-mode 1)
+   (hs-hide-level 2)
+   (should
+    (string=
+     (python-tests-visible-string)
+     "
+def f():
+    if 0:
+
+def g():
+    pass
+"))))
+
+(ert-deftest python-hideshow-hide-all-1 ()
+  "Should hide all blocks."
+  (python-tests-with-temp-buffer
+   "if 0:
+
+    aaa
+    l = [i for i in range(5)
+         if i < 3]
+    ccc
+    abc = o.match(1, 2, 3)
+    ddd
+
+def f():
+    pass
+"
+   (hs-minor-mode 1)
+   (hs-hide-all)
+   (should
+    (string=
+     (python-tests-visible-string)
+     "if 0:
+
+def f():
+"))))
+
+(ert-deftest python-hideshow-hide-all-2 ()
+  "Should hide comments."
+  (python-tests-with-temp-buffer
+   "
+# Multi line
+# comment
+
+\"\"\"
+# Multi line
+# string
+\"\"\"
+"
+   (hs-minor-mode 1)
+   (hs-hide-all)
+   (should
+    (string=
+     (python-tests-visible-string)
+     "
+# Multi line
+
+\"\"\"
+# Multi line
+# string
+\"\"\"
+"))))
+
+(ert-deftest python-hideshow-hide-all-3 ()
+  "Should not hide comments when `hs-hide-comments-when-hiding-all' is nil."
+  (python-tests-with-temp-buffer
+   "
+# Multi line
+# comment
+
+\"\"\"
+# Multi line
+# string
+\"\"\"
+"
+   (hs-minor-mode 1)
+   (let ((hs-hide-comments-when-hiding-all nil))
+     (hs-hide-all))
+   (should
+    (string=
+     (python-tests-visible-string)
+     "
+# Multi line
+# comment
+
+\"\"\"
+# Multi line
+# string
+\"\"\"
+"))))
+
+(ert-deftest python-hideshow-hide-block-1 ()
+  "Should hide current block."
+  (python-tests-with-temp-buffer
+   "
+if 0:
+
+    aaa
+    l = [i for i in range(5)
+         if i < 3]
+    ccc
+    abc = o.match(1, 2, 3)
+    ddd
+
+def f():
+    pass
+"
+   (hs-minor-mode 1)
+   (python-tests-look-at "ddd")
+   (forward-line)
+   (hs-hide-block)
+   (should
+    (string=
+     (python-tests-visible-string)
+     "
+if 0:
+
+def f():
+    pass
+"))))
+
 
 (ert-deftest python-tests--python-nav-end-of-statement--infloop ()
   "Checks that `python-nav-end-of-statement' doesn't infloop in a