Handle folding of nested defs correctly

wilywampa · wilywampa · commit d113fe0a3030 · 2014-11-17T14:01:49.000-07:00
diff --git a/autoload/pymode/folding.vim b/autoload/pymode/folding.vim
@@ -57,6 +57,21 @@ fun! pymode#folding#expr(lnum) "{{{
         return "<".(indent / &shiftwidth + 1)
     endif
 
+    " Handle nested defs
+    let last_def = pymode#motion#BlockStart(a:lnum, s:def_regex)
+    if getline(last_def) =~ s:def_regex
+        let last_def_end = pymode#motion#BlockEnd(last_def)
+        if last_def_end < line('$')
+            let nested = getline(pymode#motion#BlockStart(last_def - 1)) =~ s:def_regex
+            if nested && getline(nextnonblank(a:lnum)) !~ s:def_regex
+                let fold_end = min([prevnonblank(last_def_end - 1) + 2, last_def_end])
+                if a:lnum == fold_end
+                    return 's1'
+                endif
+            endif
+        endif
+    endif
+
     if line =~ s:blank_regex
         if prev_line =~ s:blank_regex
             if indent(a:lnum + 1) == 0 && getline(a:lnum + 1) !~ s:blank_regex
diff --git a/autoload/pymode/motion.vim b/autoload/pymode/motion.vim
@@ -32,15 +32,15 @@ endfunction "}}}
 fun! pymode#motion#select(pattern, inner) "{{{
     let cnt = v:count1 - 1
     let orig = getpos('.')[1:2]
-    let snum = s:BlockStart(orig[0], a:pattern)
+    let snum = pymode#motion#BlockStart(orig[0], a:pattern)
     if getline(snum) !~ a:pattern
         return 0
     endif
-    let enum = s:BlockEnd(snum, indent(snum))
+    let enum = pymode#motion#BlockEnd(snum, indent(snum))
     while cnt
         let lnum = search(a:pattern, 'nW')
         if lnum
-            let enum = s:BlockEnd(lnum, indent(lnum))
+            let enum = pymode#motion#BlockEnd(lnum, indent(lnum))
             call cursor(enum, 1)
         endif
         let cnt = cnt - 1
@@ -58,7 +58,7 @@ fun! pymode#motion#select(pattern, inner) "{{{
 endfunction "}}}
 
 
-fun! s:BlockStart(lnum, ...) "{{{
+fun! pymode#motion#BlockStart(lnum, ...) "{{{
     let pattern = a:0 ? a:1 : '^\s*\(@\|class\s.*:\|def\s\)'
     let lnum = a:lnum + 1
     let indent = 100
@@ -82,7 +82,7 @@ fun! s:BlockStart(lnum, ...) "{{{
 endfunction "}}}
 
 
-fun! s:BlockEnd(lnum, ...) "{{{
+fun! pymode#motion#BlockEnd(lnum, ...) "{{{
     let indent = a:0 ? a:1 : indent(a:lnum)
     let lnum = a:lnum
     while lnum
