Make folding nested defs much faster

wilywampa · wilywampa · commit 906eff560e0a · 2014-11-17T18:52:46.000-07:00
diff --git a/autoload/pymode/folding.vim b/autoload/pymode/folding.vim
@@ -58,18 +58,43 @@ fun! pymode#folding#expr(lnum) "{{{
     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'
+    if indent(prevnonblank(a:lnum))
+        let curpos = getcurpos()
+        try
+            let last_block = s:BlockStart(a:lnum)
+            let last_block_indent = indent(last_block)
+
+            " Check if last class/def is not indented and therefore can't be
+            " nested and make sure it is a class/def block instead of a zero
+            " indented regular statement
+            if last_block_indent && getline(last_block) =~ s:def_regex
+                " Note: This relies on the cursor position being set by s:BlockStart
+                let next_def = searchpos('^\s*def \w', 'nW')[0]
+                let next_def_indent = next_def ? indent(next_def) : -1
+                let last_block_end = s:BlockEnd(last_block)
+
+                " If the next def has the same or greater indent than the
+                " previous def, it is either nested at the same level or
+                " nested one level deeper, and in either case will have its
+                " own fold. If the class/def containing the current line is on
+                " the first line it can't be nested, and if the this block
+                " ends on the last line, it contains no trailing code that
+                " should not be folded. Otherwise, we know the current line
+                " is at the end of a nested def.
+                if next_def_indent < last_block_indent && last_block > 1 && last_block_end < line('$')
+
+                    " Include up to one blank line in the fold
+                    let fold_end = min([prevnonblank(last_block_end - 1) + 1, last_block_end])
+                    if a:lnum == fold_end
+                        return 's1'
+                    else
+                        return '='
+                    endif
                 endif
             endif
-        endif
+        finally
+            call setpos('.', curpos)
+        endtry
     endif
 
     if line =~ s:blank_regex
@@ -91,5 +116,17 @@ fun! pymode#folding#expr(lnum) "{{{
 
 endfunction "}}}
 
+fun! s:BlockStart(lnum) "{{{
+    " Note: Make sure to reset cursor position after using this function.
+    call cursor(a:lnum, 0)
+    let max_indent = max([indent(prevnonblank(a:lnum)) - &shiftwidth, 0])
+    return searchpos('\v^(\s{,'.max_indent.'}(def |class |\@)\w|[^ \t#])', 'bcnW')[0]
+endfunction "}}}
+
+fun! s:BlockEnd(lnum) "{{{
+    " Note: Make sure to reset cursor position after using this function.
+    call cursor(a:lnum, 0)
+    return searchpos('\v^\s{,'.indent('.').'}\S', 'nW')[0] - 1
+endfunction "}}}
 
 " vim: fdm=marker:fdl=0
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 = pymode#motion#BlockStart(orig[0], a:pattern)
+    let snum = s:BlockStart(orig[0], a:pattern)
     if getline(snum) !~ a:pattern
         return 0
     endif
-    let enum = pymode#motion#BlockEnd(snum, indent(snum))
+    let enum = s:BlockEnd(snum, indent(snum))
     while cnt
         let lnum = search(a:pattern, 'nW')
         if lnum
-            let enum = pymode#motion#BlockEnd(lnum, indent(lnum))
+            let enum = s: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! pymode#motion#BlockStart(lnum, ...) "{{{
+fun! s: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! pymode#motion#BlockStart(lnum, ...) "{{{
 endfunction "}}}
 
 
-fun! pymode#motion#BlockEnd(lnum, ...) "{{{
+fun! s:BlockEnd(lnum, ...) "{{{
     let indent = a:0 ? a:1 : indent(a:lnum)
     let lnum = a:lnum
     while lnum
