Merge pull request python-mode#207 from alvinfrancis/develop

klen · klen · commit 04948cfef885 · 2013-04-26T01:28:03.000-07:00
Modify folding
diff --git a/autoload/pymode/folding.vim b/autoload/pymode/folding.vim
@@ -2,12 +2,16 @@
 
 
 let s:blank_regex = '^\s*$'
-let s:def_regex = '^\s*\(class\|def\) \w\+'
+let s:def_regex = '^\s*\%(class\|def\) \w\+'
+let s:decorator_regex = '^\s*@'
+let s:doc_begin_regex = '^\s*\%("""\|''''''\)'
+let s:doc_end_regex = '\%("""\|''''''\)\s*$'
+let s:doc_line_regex = '^\s*\("""\|''''''\).\+\1\s*$'
 
 
 fun! pymode#folding#text() " {{{
     let fs = v:foldstart
-    while getline(fs) =~ '^\s*@' 
+    while getline(fs) =~ '\%(^\s*@\)\|\%(^\s*\%("""\|''''''\)\s*$\)'
         let fs = nextnonblank(fs + 1)
     endwhile
     let line = getline(fs)
@@ -21,6 +25,7 @@ fun! pymode#folding#text() " {{{
     let line = substitute(line, '\t', onetab, 'g')
 
     let line = strpart(line, 0, windowwidth - 2 -len(foldedlinecount))
+    let line = substitute(line, '\%("""\|''''''\)', '', '')
     let fillcharcount = windowwidth - len(line) - len(foldedlinecount)
     return line . '…' . repeat(" ",fillcharcount) . foldedlinecount . '…' . ' '
 endfunction "}}}
@@ -30,21 +35,32 @@ fun! pymode#folding#expr(lnum) "{{{
 
     let line = getline(a:lnum)
     let indent = indent(a:lnum)
+	let prev_line = getline(a:lnum - 1)
+
+    if line =~ s:def_regex || line =~ s:decorator_regex
+		if prev_line =~ s:decorator_regex
+			return '='
+		else
+			return ">".(indent / &shiftwidth + 1)
+		endif
+    endif
 
-    if line =~ s:def_regex
-        return ">".(indent / &shiftwidth + 1)
+    if line =~ s:doc_begin_regex
+				\ && line !~ s:doc_line_regex
+				\ && prev_line =~ s:def_regex
+		return ">".(indent / &shiftwidth + 1)
     endif
 
-    if line =~ '^\s*@'
-        return -1
+    if line =~ s:doc_end_regex
+				\ && line !~ s:doc_line_regex
+		return "<".(indent / &shiftwidth + 1)
     endif
 
     if line =~ s:blank_regex
-        let prev_line = getline(a:lnum - 1)
         if prev_line =~ s:blank_regex
             return -1
         else
-            return foldlevel(prevnonblank(a:lnum))
+            return '='
         endif
     endif
 
