Improve Python folding behaviour

after/ftplugin/python.vim

@@ -28,3 +28,46 @@ let python_highlight_all=1
 
 " Mappings
 nnoremap K :YcmCompleter GetDoc<CR>
+
+" Set custom fold expression
+setlocal foldmethod=expr
+setlocal foldexpr=PythonFold(v:lnum)
+
+" Custom fold function greedyily matches following blank lines
+function! PythonFold(lnum)
+  let l:line = getline(a:lnum)
+
+  " Blank lines, comments, and docstrings use previous fold level
+  if l:line =~ '^\(\s*\|#.*\|"\(""\)\=.*\)$'
+    return '='
+  " Keywords beginning indented blocks start a fold
+  elseif l:line =~ '^\s*class\s\+\w*(\w*)\s*:\s*$'
+    \ || l:line =~ '^\s*def\s\+\w\+\s*(.*)\s*:\s*$'
+    \ || l:line =~ '^\s*if\s\+.*:\s*$'
+    \ || l:line =~ '^\s*elif\s\+.*:\s*$'
+    \ || l:line =~ '^\s*else\s*:\s*$'
+    \ || l:line =~ '^\s*for\s\+.*:\s*$'
+    \ || l:line =~ '^\s*while\s\+.*:\s*$'
+    \ || l:line =~ '^\s*try\s*:\s*$'
+    \ || l:line =~ '^\s*except\s*.*:\s*$'
+    return '>'.string((indent(a:lnum) / &shiftwidth) + 1)
+  " Opening curly braces, not in a string, add a fold level
+  elseif l:line =~ '{' && l:line !~ '}' && l:line !~ "'.*{.*'" && l:line !~ '".*{.*"'
+    return 'a1'
+  " Closing curly braces, not in a string, substract a fold level
+  elseif l:line =~ '}' && l:line !~ '{' && l:line !~ "'.*}.*'" && l:line !~ '".*}.*"'
+    return 's1'
+  " Opening square braces, not in a string, add a fold level
+  elseif l:line =~ '[' && l:line !~ ']' && l:line !~ "'.*[.*'" && l:line !~ '".*].*"'
+    return 'a1'
+  " Closing square braces, not in a string, substract a fold level
+  elseif l:line =~ ']' && l:line !~ '[' && l:line !~ "'.*].*'" && l:line !~ '".*[.*"'
+    return 's1'
+  " Calculate lines with a lower indent than the previous line
+  elseif indent(a:lnum) < indent(a:lnum - 1)
+    return string((indent(a:lnum) / &shiftwidth))
+  endif
+
+  " Finally all unmatched lines use fold level from previous line
+  return '='
+endfunction