python-mode · fmv1992 · Feb 6, 2017 · Nov 18, 2016 · Nov 18, 2016 · Feb 6, 2017
diff --git a/autoload/pymode/folding.vim b/autoload/pymode/folding.vim
@@ -1,16 +1,25 @@
 " Python-mode folding functions
 
+" Notice that folding is based on single line so complex regular expressions
+" that take previous line into consideration are not fit for the job.
 
+" Regex definitions for correct folding
 let s:def_regex = g:pymode_folding_regex
 let s:blank_regex = '^\s*$'
-let s:decorator_regex = '^\s*@'
-let s:doc_begin_regex = '^\s*\%("""\|''''''\)'
+" Spyder, a very popular IDE for python has a template which includes
+" '@author:' ; thus the regex below.
+let s:decorator_regex = '^\s*@\(author:\)\@!' 
+let s:doc_begin_regex = '^\s*[uU]\=\%("""\|''''''\)'
 let s:doc_end_regex = '\%("""\|''''''\)\s*$'
-let s:doc_line_regex = '^\s*\("""\|''''''\).\+\1\s*$'
+" This one is needed for the while loop to count for opening and closing
+" docstrings.
+let s:doc_general_regex = '\%("""\|''''''\)'
+let s:doc_line_regex = '^\s*[uU]\=\("""\|''''''\).\+\1\s*$'
 let s:symbol = matchstr(&fillchars, 'fold:\zs.')  " handles multibyte characters
 if s:symbol == ''
     let s:symbol = ' '
 endif
+" ''''''''
 
 
 fun! pymode#folding#text() " {{{
@@ -33,24 +42,29 @@ 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 line = substitute(line, '[uU]\=\%("""\|''''''\)', '', '')
     let fillcharcount = windowwidth - len(line) - len(foldedlinecount) + 1
     return line . ' ' . repeat(s:symbol, fillcharcount) . ' ' . foldedlinecount
 endfunction "}}}
 
-
 fun! pymode#folding#expr(lnum) "{{{
 
     let line = getline(a:lnum)
     let indent = indent(a:lnum)
     let prev_line = getline(a:lnum - 1)
+    let next_line = getline(a:lnum + 1)
 
+    " Decorators {{{
     if line =~ s:decorator_regex
         return ">".(indent / &shiftwidth + 1)
-    endif
+    endif "}}}
 
+    " Definition {{{
     if line =~ s:def_regex
-        " single line def
+        " If indent of this line is greater or equal than line below
+        " and previous non blank line does not end with : (that is, is not a
+        " definition)
+        " Keep the same indentation
         if indent(a:lnum) >= indent(a:lnum+1) && getline(prevnonblank(a:lnum)) !~ ':\s*$'
             return '='
         endif
@@ -71,16 +85,35 @@ fun! pymode#folding#expr(lnum) "{{{
         else
             return ">".(indent / &shiftwidth + 1)
         endif
-    endif
+    endif "}}}
 
-    if line =~ s:doc_begin_regex && line !~ s:doc_line_regex && prev_line =~ s:def_regex
-        return ">".(indent / &shiftwidth + 1)
+    " Docstrings {{{
+
+    " TODO: A while loop now counts the number of open and closed folding in
+    " order to determine if it is a closing or opening folding.
+    " It is working but looks like it is an overkill.
+
+    " Notice that an effect of this is that other docstring matches will not
+    " be one liners.
+    if line =~ s:doc_line_regex
+        return "="
     endif
 
-    if line =~ s:doc_end_regex && line !~ s:doc_line_regex
-        return "<".(indent / &shiftwidth + 1)
+    if line =~ s:doc_begin_regex
+            " echom 'just entering'
+        if s:Is_opening_folding(a:lnum)
+            " echom 'entering at line ' . a:lnum
+            return ">".(indent / &shiftwidth + 1)
+        endif
     endif
+    if line =~ s:doc_end_regex
+        if !s:Is_opening_folding(a:lnum)
+            " echom 'leaving at line ' . a:lnum
+            return "<".(indent / &shiftwidth + 1)
+        endif
+    endif "}}}
 
+    " Nested Definitions {{{
     " Handle nested defs but only for files shorter than
     " g:pymode_folding_nest_limit lines due to performance concerns
     if line('$') < g:pymode_folding_nest_limit && indent(prevnonblank(a:lnum))
@@ -125,18 +158,25 @@ fun! pymode#folding#expr(lnum) "{{{
         finally
             call setpos('.', curpos)
         endtry
-    endif
+    endif " }}}
 
+    " Blank Line {{{
     if line =~ s:blank_regex
         if prev_line =~ s:blank_regex
-            if indent(a:lnum + 1) == 0 && getline(a:lnum + 1) !~ s:blank_regex
-                return 0
+            if indent(a:lnum + 1) == 0 && next_line !~ s:blank_regex && next_line !~ s:doc_general_regex
+                if s:Is_opening_folding(a:lnum)
+                    " echom a:lnum
+                    return "="
+                else
+                    " echom "not " . a:lnum
+                    return 0
+                endif
             endif
             return -1
         else
             return '='
         endif
-    endif
+    endif " }}}
 
     return '='
 
@@ -174,4 +214,52 @@ fun! s:BlockEnd(lnum) "{{{
     return searchpos('\v^\s{,'.indent('.').'}\S', 'nW')[0] - 1
 endfunction "}}}
 
+function! s:Is_opening_folding(lnum) "{{{
+    " Helper function to see if docstring is opening or closing
+    let number_of_folding = 0  " To be analized if odd/even to inform if it is opening or closing.
+    let has_open_docstring = 0  " To inform is already has an open docstring.
+    let extra_docstrings = 0  " To help skipping ''' and """ which are not docstrings
+
+    " The idea of this part of the function is to identify real docstrings and
+    " not just triple quotes (that could be a regular string).
+    "
+    " Iterater over all lines from the start until current line (inclusive)
+    for i in range(1, a:lnum)
+        let i_line = getline(i)
+
+        if i_line =~ s:doc_line_regex 
+            " echom "case 00 on line " . i
+            continue
+        endif
+
+        if i_line =~ s:doc_begin_regex && ! has_open_docstring
+            " echom "case 01 on line " . i
+            " This causes the loop to continue if there is a triple quote which
+            " is not a docstring.
+            if extra_docstrings > 0
+                let extra_docstrings = extra_docstrings - 1
+                continue
+            else
+                let has_open_docstring = 1
+                let number_of_folding = number_of_folding + 1
+            endif
+        " If it is an end doc and has an open docstring.
+        elseif i_line =~ s:doc_end_regex && has_open_docstring
+            " echom "case 02 on line " . i
+            let has_open_docstring = 0
+            let number_of_folding = number_of_folding + 1
+
+        elseif i_line =~ s:doc_general_regex
+            " echom "extra docstrings on line " . i
+            let extra_docstrings = extra_docstrings + 1
+        endif 
+    endfor
+
+    if fmod(number_of_folding, 2) == 1 "If odd then it is an opening
+        return 1
+    else
+        return 0
+    endif
+endfunction "}}}
+
 " vim: fdm=marker:fdl=0
diff --git a/plugin/pymode.vim b/plugin/pymode.vim
@@ -39,7 +39,7 @@ call pymode#default("g:pymode_folding", 1)
 " Maximum file length to check for nested class/def statements
 call pymode#default("g:pymode_folding_nest_limit", 1000)
 " Change for folding customization (by example enable fold for 'if', 'for')
-call pymode#default("g:pymode_folding_regex", '^\s*\%(class\|def\|async\s\+def\) \w\+')
+call pymode#default("g:pymode_folding_regex", '^\s*\%(class\|def\|async\s\+def\) .\+\(:\s\+\w\)\@!')
 
 " Enable/disable python motion operators
 call pymode#default("g:pymode_motion", 1)