oxml: add CT_Body.add_p()

Steve Canny · Steve Canny · commit 8e7a17ae79b7 · 2013-12-20T20:10:44.000-08:00
diff --git a/docx/oxml/__init__.py b/docx/oxml/__init__.py
@@ -13,11 +13,13 @@
 """
 
 from docx.oxml.base import register_custom_element_class
+from docx.oxml.parts import CT_Body
 from docx.oxml.text import CT_P
 
 
 # ===========================================================================
 # custom element class mappings
 # ===========================================================================
 
+register_custom_element_class('w:body',     CT_Body)
 register_custom_element_class('w:p',        CT_P)
diff --git a/docx/oxml/parts.py b/docx/oxml/parts.py
@@ -0,0 +1,41 @@
+# -*- coding: utf-8 -*-
+#
+# oxml/parts.py
+#
+# Copyright (C) 2013 Steve Canny scanny@cisco.com
+#
+# This module is part of python-docx and is released under the MIT License:
+# http://www.opensource.org/licenses/mit-license.php
+
+"""
+Custom element classes that correspond to OPC parts like <w:document>.
+"""
+
+from docx.oxml.base import OxmlBaseElement
+from docx.oxml.text import CT_P
+
+
+class CT_Body(OxmlBaseElement):
+    """
+    ``<w:body>``, the container element for the main document story in
+    ``document.xml``.
+    """
+    def add_p(self):
+        """
+        Return a new <w:p> element that has been added at the end of any
+        existing body content.
+        """
+        p = CT_P.new()
+        if hasattr(self, 'sectPr'):
+            self.sectPr.addprevious(p)
+        else:
+            self.append(p)
+        return p
+
+    @property
+    def _has_sectPr(self):
+        """
+        Return True if this <w:body> element has a <w:sectPr> child element,
+        False otherwise.
+        """
+        return hasattr(self, 'sectPr')
diff --git a/tests/oxml/test_parts.py b/tests/oxml/test_parts.py
@@ -0,0 +1,34 @@
+# -*- coding: utf-8 -*-
+#
+# test_parts.py
+#
+# Copyright (C) 2013 Steve Canny scanny@cisco.com
+#
+# This module is part of python-docx and is released under the MIT License:
+# http://www.opensource.org/licenses/mit-license.php
+
+"""Test suite for the docx.oxml.parts module."""
+
+from docx.oxml.text import CT_P
+
+from ..unitdata import a_body
+
+
+class DescribeCT_Body(object):
+
+    def it_can_add_a_p_to_itself(self):
+        """
+        Return a newly created |CT_P| element that has been added after any
+        existing content.
+        """
+        cases = (
+            (a_body(),               a_body().with_p()),
+            (a_body().with_sectPr(), a_body().with_p().with_sectPr()),
+        )
+        for before_body_bldr, after_body_bldr in cases:
+            body = before_body_bldr.element
+            # exercise -----------------
+            p = body.add_p()
+            # verify -------------------
+            assert body.xml == after_body_bldr.xml
+            assert isinstance(p, CT_P)
diff --git a/tests/unitdata.py b/tests/unitdata.py
@@ -43,6 +43,47 @@ def with_nsdecls(self):
         return self
 
 
+class CT_BodyBuilder(BaseBuilder):
+    """
+    Test data builder for CT_Body (<w:body>) XML element that appears in
+    document.xml files.
+    """
+    def __init__(self):
+        """Establish instance variables with default values"""
+        super(CT_BodyBuilder, self).__init__()
+        self._p = None
+        self._sectPr = None
+
+    @property
+    def is_empty(self):
+        return self._p is None and self._sectPr is None
+
+    def with_p(self):
+        """Add an empty paragraph element"""
+        self._p = a_p()
+        return self
+
+    def with_sectPr(self):
+        """Add an empty section properties element"""
+        self._sectPr = a_sectPr()
+        return self
+
+    @property
+    def xml(self):
+        """Return element XML based on attribute settings"""
+        indent = ' ' * self._indent
+        if self.is_empty:
+            xml = '%s<w:body %s/>\n' % (indent, nsdecls('w'))
+        else:
+            xml = '%s<w:body %s>\n' % (indent, nsdecls('w'))
+            if self._p:
+                xml += self._p.with_indent(self._indent+2).xml
+            if self._sectPr:
+                xml += self._sectPr.with_indent(self._indent+2).xml
+            xml += '%s</w:body>\n' % indent
+        return xml
+
+
 class CT_PBuilder(BaseBuilder):
     """
     Test data builder for a CT_P (<w:p>) XML element that appears within the
@@ -60,6 +101,33 @@ def xml(self):
         return xml
 
 
+class CT_SectPrBuilder(BaseBuilder):
+    """
+    Test data builder for a CT_SectPr (<w:sectPr>) XML element that appears
+    within the body element of a document.xml file.
+    """
+    def __init__(self):
+        """Establish instance variables with default values"""
+        super(CT_SectPrBuilder, self).__init__()
+
+    @property
+    def xml(self):
+        """Return element XML based on attribute settings"""
+        tmpl = '%s<w:sectPr/>\n'
+        indent = ' ' * self._indent
+        return tmpl % (indent)
+
+
+def a_body():
+    """Return a CT_BodyBuilder instance"""
+    return CT_BodyBuilder()
+
+
 def a_p():
     """Return a CT_PBuilder instance"""
     return CT_PBuilder()
+
+
+def a_sectPr():
+    """Return a CT_SectPr instance"""
+    return CT_SectPrBuilder()
