Handle heading levels 1 to 6 in markdown writer

1 files changed, 30 insertions, 4 deletions

diff --git a/mammoth/writers/markdown.py b/mammoth/writers/markdown.py
index 604b18a..4a2476d 100644
--- a/mammoth/writers/markdown.py
+++ b/mammoth/writers/markdown.py
@@ -3,20 +3,46 @@ from __future__ import unicode_literals
 import re
 
 
+
+class _Wrapped(object):
+    def __init__(self, start, end):
+        self._start = start
+        self._end = end
+    
+    def __call__(self):
+        return self._start, self._end
+
+
+def _init_writers():
+    writers = {
+        "p": _Wrapped("", "\n\n"),
+    }
+    
+    for level in range(1, 7):
+        writers["h{0}".format(level)] = _Wrapped("#" * level + " ", "\n\n")
+    
+    return writers
+
+
+_writers = _init_writers()
+
+
 class MarkdownWriter(object):
     def __init__(self):
         self._fragments = []
+        self._element_stack = []
     
     def text(self, text):
         self._fragments.append(_escape_markdown(text))
     
     def start(self, name, attributes=None):
-        if name == "h1":
-            self._fragments.append("# ")
+        start, end = _writers[name]()
+        self._fragments.append(start)
+        self._element_stack.append(end)
 
     def end(self, name):
-        if name in ["p", "h1"]:
-            self._fragments.append("\n\n")
+        end = self._element_stack.pop()
+        self._fragments.append(end)
     
     def self_closing(self, name, attributes=None):
         pass