diff --git a/ftplugin/python/SimpylFold.vim b/ftplugin/python/SimpylFold.vim index 00fa32d..69db1b6 100644 --- a/ftplugin/python/SimpylFold.vim +++ b/ftplugin/python/SimpylFold.vim @@ -5,6 +5,9 @@ let b:loaded_SimpylFold = 1 let s:blank_regex = '^\s*$' let s:def_regex = '^\s*\%(class\|def\) \w\+' +let s:docstring_start_regex = '^\s*\("""\|''''''\)\%(.*\1\s*$\)\@!' +let s:docstring_end_single_regex = '''''''\s*$' +let s:docstring_end_double_regex = '"""\s*$' " Determine the number of containing class or function definitions for the " given line @@ -66,6 +69,7 @@ function! SimpylFold(lnum) " the cache of results of calls to `s:NumContainingDefs` if !exists('b:last_folded_line') || b:last_folded_line > a:lnum let b:cache_NumContainingDefs = {} + let b:in_docstring = 0 endif let b:last_folded_line = a:lnum @@ -84,10 +88,26 @@ function! SimpylFold(lnum) endif endif - " Otherwise, its fold level is equal to its number of containing - " definitions, plus 1, if this line starts a definition of its own - let this_fl = s:NumContainingDefs(a:lnum) + (line =~ s:def_regex) + let docstring_match = matchlist(line, s:docstring_start_regex) + if !b:in_docstring && len(docstring_match) + let this_fl = s:NumContainingDefs(a:lnum) + 1 + let b:in_docstring = 1 + if docstring_match[1] == '"""' + let b:docstring_end_regex = s:docstring_end_double_regex + else + let b:docstring_end_regex = s:docstring_end_single_regex + endif + elseif b:in_docstring + let this_fl = s:NumContainingDefs(a:lnum) + 1 + if line =~ b:docstring_end_regex + let b:in_docstring = 0 + endif + else + " Otherwise, its fold level is equal to its number of containing + " definitions, plus 1, if this line starts a definition of its own + let this_fl = s:NumContainingDefs(a:lnum) + (line =~ s:def_regex) + endif " If the very next line starts a definition with the same fold level as " this one, explicitly indicate that a fold ends here if getline(a:lnum + 1) =~ s:def_regex && SimpylFold(a:lnum + 1) == this_fl