From d446595f9e690e9f50eeda0f0e144320ccca72dc Mon Sep 17 00:00:00 2001
From: Kit La Touche <kit@simpleenergy.com>
Date: Tue, 12 May 2015 11:17:15 -0600
Subject: [PATCH] Add support for python-mode style folding

As a python-mode user in my day-to-day work, I would like the folding
styles of vim-markdown and python-mode to match, so that there's less
cognitive overhead in switching between my two most frequently-used
contexts.

To this end, I've made a distinct set of folding logic, keyed off of the
g:vim_markdown_folding_style_pythonic variable, so that using this style
is opt-in.
---
 README.md                   |  9 ++++
 after/ftplugin/markdown.vim | 97 ++++++++++++++++++++++++++-----------
 test/python-folding.vader   | 39 +++++++++++++++
 3 files changed, 116 insertions(+), 29 deletions(-)
 create mode 100644 test/python-folding.vader

diff --git a/README.md b/README.md
index 5ddf6ed..925fd67 100644
--- a/README.md
+++ b/README.md
@@ -70,6 +70,15 @@ To enable/disable folding use Vim's standard folding configuration.
 set [no]foldenable
 ```
 
+### Change fold style
+
+To fold in a style like [python-mode](https://github.com/klen/python-mode), add
+the following to your `.vimrc`:
+
+```vim
+let g:vim_markdown_folding_style_pythonic = 1
+```
+
 ### Disable Default Key Mappings
 
 Add the following line to your `.vimrc` to disable default key mappings:
diff --git a/after/ftplugin/markdown.vim b/after/ftplugin/markdown.vim
index a140c25..e635f6d 100644
--- a/after/ftplugin/markdown.vim
+++ b/after/ftplugin/markdown.vim
@@ -15,38 +15,77 @@ func! s:effective_line(lnum)
     return (line !~ '^[=-#]' || s:is_mkdCode(a:lnum)) ? '' : line
 endfunc
 
-func! Foldexpr_markdown(lnum)
-    let l2 = s:effective_line(a:lnum+1)
-    if  l2 =~ '^==\+\s*'
-        " next line is underlined (level 1)
-        return '>1'
-    elseif l2 =~ '^--\+\s*'
-        " next line is underlined (level 2)
-        return '>2'
-    endif
+if get(g:, "vim_markdown_folding_style_pythonic", 0)
+    func! Foldexpr_markdown(lnum)
+        let l2 = s:effective_line(a:lnum+1)
+        if  l2 =~ '^==\+\s*'
+            " next line is underlined (level 1)
+            return '>0'
+        elseif l2 =~ '^--\+\s*'
+            " next line is underlined (level 2)
+            return '>1'
+        endif
 
-    let l1 = s:effective_line(a:lnum)
-    if l1 =~ '^#'
-        " don't include the section title in the fold
-        return '-1'
-    endif
+        let l1 = s:effective_line(a:lnum)
+        if l1 =~ '^#'
+            " current line starts with hashes
+            return '>'.(matchend(l1, '^#\+') - 1)
+        elseif a:lnum == 1
+            " fold any 'preamble'
+            return '>1'
+        else
+            " keep previous foldlevel
+            return '='
+        endif
+    endfunc
 
-    if (a:lnum == 1)
-        let l0 = ''
-    else
-        let l0 = s:effective_line(a:lnum-1)
-    endif
-    if l0 =~ '^#'
-        " current line starts with hashes
-        return '>'.matchend(l0, '^#\+')
-    else
-        " keep previous foldlevel
-        return '='
-    endif
-endfunc
+    fun! Foldtext_markdown()
+        let line = getline(v:foldstart)
+        let has_numbers = &number || &relativenumber
+        let nucolwidth = &fdc + has_numbers * &numberwidth
+        let windowwidth = winwidth(0) - nucolwidth - 6
+        let foldedlinecount = v:foldend - v:foldstart
+        let line = strpart(line, 0, windowwidth - 2 -len(foldedlinecount))
+        let line = substitute(line, '\%("""\|''''''\)', '', '')
+        let fillcharcount = windowwidth - len(line) - len(foldedlinecount) + 1
+        return line . ' ' . repeat("-", fillcharcount) . ' ' . foldedlinecount
+    endfunc
+else
+    func! Foldexpr_markdown(lnum)
+        let l2 = s:effective_line(a:lnum+1)
+        if  l2 =~ '^==\+\s*'
+            " next line is underlined (level 1)
+            return '>1'
+        elseif l2 =~ '^--\+\s*'
+            " next line is underlined (level 2)
+            return '>2'
+        endif
+
+        let l1 = s:effective_line(a:lnum)
+        if l1 =~ '^#'
+            " don't include the section title in the fold
+            return '-1'
+        endif
 
+        if (a:lnum == 1)
+            let l0 = ''
+        else
+            let l0 = s:effective_line(a:lnum-1)
+        endif
+        if l0 =~ '^#'
+            " current line starts with hashes
+            return '>'.matchend(l0, '^#\+')
+        else
+            " keep previous foldlevel
+            return '='
+        endif
+    endfunc
+endif
 
 if !get(g:, "vim_markdown_folding_disabled", 0)
-  setlocal foldexpr=Foldexpr_markdown(v:lnum)
-  setlocal foldmethod=expr
+    setlocal foldexpr=Foldexpr_markdown(v:lnum)
+    setlocal foldmethod=expr
+    if get(g:, "vim_markdown_folding_style_pythonic", 0)
+        setlocal foldtext=Foldtext_markdown()
+    endif
 endif
diff --git a/test/python-folding.vader b/test/python-folding.vader
new file mode 100644
index 0000000..55d46de
--- /dev/null
+++ b/test/python-folding.vader
@@ -0,0 +1,39 @@
+Before:
+  let g:vim_markdown_folding_style_pythonic = 1
+  source ../after/ftplugin/markdown.vim
+
+After:
+  setlocal foldexpr=0
+  setlocal foldmethod=manual
+
+Given markdown;
+# Title
+
+## Chapter 1
+
+```
+This is code block
+# This is just a comment
+```
+
+## Chapter 2
+
+foobar
+
+Execute (fold level # in code block):
+  AssertEqual foldlevel(1), 0, '# Title'
+  AssertEqual foldlevel(3), 1, '## Chapter 1'
+  AssertEqual foldlevel(7), 1, '# This is just a comment'
+  AssertEqual foldlevel(8), 1, '```'
+  AssertEqual foldlevel(10), 1, '## Chapter 2'
+  AssertEqual foldlevel(12), 1, 'foobar'
+
+Given markdown;
+
+==+ Fold Level 1
+
+--+ Fold Level 2
+
+Execute (fold level ==+, --+):
+  AssertEqual foldlevel(2), 0, '==+'
+  AssertEqual foldlevel(4), 1, '--+'
-- 
2.39.5