Handle heading levels 1 to 6 in markdown writer

2 files changed, 39 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
diff --git a/tests/writers/markdown_tests.py b/tests/writers/markdown_tests.py
index f334dfa..e1cdf77 100644
--- a/tests/writers/markdown_tests.py
+++ b/tests/writers/markdown_tests.py
@@ -28,6 +28,15 @@ def h1_elements_are_converted_to_heading_with_leading_hash():
     writer.text("Hello");
     writer.end("h1");
     assert_equal("# Hello\n\n", writer.as_string())
+
+
+@istest
+def h6_elements_are_converted_to_heading_with_six_leading_hashes():
+    writer = _create_writer()
+    writer.start("h6");
+    writer.text("Hello");
+    writer.end("h6");
+    assert_equal("###### Hello\n\n", writer.as_string())