doc: add Document.add_heading()

Steve Canny · Steve Canny · commit 1d19d4589bcb · 2015-02-14T17:30:47.000-08:00
diff --git a/docx/document.py b/docx/document.py
@@ -24,6 +24,20 @@ def __init__(self, element, part):
         self._part = part
         self.__body = None
 
+    def add_heading(self, text='', level=1):
+        """
+        Return a heading paragraph newly added to the end of the document,
+        containing *text* and having its paragraph style determined by
+        *level*. If *level* is 0, the style is set to `Title`. If *level* is
+        1 (or omitted), `Heading 1` is used. Otherwise the style is set to
+        `Heading {level}`. Raises |ValueError| if *level* is outside the
+        range 0-9.
+        """
+        if not 0 <= level <= 9:
+            raise ValueError("level must be in range 0-9, got %d" % level)
+        style = 'Title' if level == 0 else 'Heading %d' % level
+        return self.add_paragraph(text, style)
+
     def add_paragraph(self, text='', style=None):
         """
         Return a paragraph newly added to the end of the document, populated
diff --git a/tests/test_document.py b/tests/test_document.py
@@ -15,11 +15,26 @@
 from docx.text.paragraph import Paragraph
 
 from .unitutil.cxml import element
-from .unitutil.mock import class_mock, instance_mock, property_mock
+from .unitutil.mock import (
+    class_mock, instance_mock, method_mock, property_mock
+)
 
 
 class DescribeDocument(object):
 
+    def it_can_add_a_heading(self, add_heading_fixture):
+        document, text, level, style, paragraph_ = add_heading_fixture
+        paragraph = document.add_heading(text, level)
+        document.add_paragraph.assert_called_once_with(text, style)
+        assert paragraph is paragraph_
+
+    def it_raises_on_heading_level_out_of_range(self):
+        document = Document(None, None)
+        with pytest.raises(ValueError):
+            document.add_heading(level=-1)
+        with pytest.raises(ValueError):
+            document.add_heading(level=10)
+
     def it_can_add_a_paragraph(self, add_paragraph_fixture):
         document, text, style, paragraph_ = add_paragraph_fixture
         paragraph = document.add_paragraph(text, style)
@@ -38,6 +53,19 @@ def it_provides_access_to_the_document_body(self, body_fixture):
 
     # fixtures -------------------------------------------------------
 
+    @pytest.fixture(params=[
+        (0, 'Title'),
+        (1, 'Heading 1'),
+        (2, 'Heading 2'),
+        (9, 'Heading 9'),
+    ])
+    def add_heading_fixture(self, request, add_paragraph_, paragraph_):
+        level, style = request.param
+        document = Document(None, None)
+        text = 'Spam vs. Bacon'
+        add_paragraph_.return_value = paragraph_
+        return document, text, level, style, paragraph_
+
     @pytest.fixture(params=[
         ('',         None),
         ('',         'Heading 1'),
@@ -63,6 +91,10 @@ def part_fixture(self, document_part_):
 
     # fixture components ---------------------------------------------
 
+    @pytest.fixture
+    def add_paragraph_(self, request):
+        return method_mock(request, Document, 'add_paragraph')
+
     @pytest.fixture
     def _Body_(self, request, body_):
         return class_mock(request, 'docx.document._Body', return_value=body_)
