Fold up to the fold start, instead of the indent start - vim

Question

I use foldmethod=indent and when I fold code like this:
def cake():
#cake!
print( "cake" )
print( "for" )
print( "you" )
I see
def cake():
#cake!
print( "cake" ) +++ 3 lines folded
but I want to see
def cake(): +++ 5 lines folded
Is there a way to do fold up to the first line (def cake():) like this?

Answer 1

Chapters 48 and 49 of Learn Vimscript the Hard Way talk about how to do that, using foldmethod=expr instead of indent. Basically you need to make a custom ftplugin and put a folding script in it; the script contains functions used to determine what fold level different lines should have.
As luck would have it, the example code given in those two chapters is for the Potion language which, like Python, is whitespace-sensitive, so it should be pretty easy to adapt it to Python. Since Vim already comes with a Python ftplugin, I think you can put the folding script described on the site into .vim/after/ftplugin/python instead of .vim/ftplugin/potion.

Answer 2

I solved this using this tutorial.
This is the finished bunch of functions:
fu! Indent_level(lnum)
return indent(a:lnum) / &shiftwidth
endfunction
fu! Next_non_blank_line(lnum)
let numlines = line('$')
let current = a:lnum + 1
while current <= numlines
if getline(current) =~? '\v\S'
return current
endif
let current += 1
endwhile
return -2
endfunction
fu! Custom_fold_expr(lnum)
if getline(a:lnum) =~? '\v^\s*$'
return '-1'
endif
let this_indent = Indent_level(a:lnum)
let next_indent = Indent_level(Next_non_blank_line(a:lnum))
if next_indent == this_indent
return this_indent
elseif next_indent < this_indent
return this_indent
elseif next_indent > this_indent
return '>' . next_indent
endif
endf
set foldexpr=Custom_fold_expr(v:lnum)
foldmethod=expr
Please don't edit the indentation of the "end" markers on this post, it looks gorgeous after you put this in your vimrc.