Wrap options based on enum

Allows wrapSquare bothSides, wrapTopAndBottom
SweetProcess · jtrain · Apr 28, 2017 · Apr 28, 2017 · Apr 28, 2017 · May 1, 2017
commit 275ff9cb42787382fd492cbf82fef32a2623a159
diff --git a/docx/enum/shape.py b/docx/enum/shape.py
@@ -6,6 +6,8 @@
 
 from __future__ import absolute_import, print_function, unicode_literals
 
+from .base import Enumeration, EnumMember
+
 
 class WD_INLINE_SHAPE_TYPE(object):
     """
@@ -35,3 +37,20 @@ class WD_ANCHOR_SHAPE_TYPE(object):
 
 
 WD_ANCHOR_SHAPE = WD_ANCHOR_SHAPE_TYPE
+
+
+class WRAP_SHAPE_TYPE(Enumeration):
+
+    __ms_name__ = ''
+    __members__ = (
+
+        EnumMember(
+            'wrapSquareBothSides', 'bothSides',
+            'A square wrapped shape with text wrapping on both sides'
+        ),
+
+        EnumMember(
+            'wrapTopAndBottom', '',
+            'A square on its own line cleared on left and right'
+        )
+    )
diff --git a/docx/oxml/__init__.py b/docx/oxml/__init__.py
@@ -96,23 +96,25 @@ def OxmlElement(nsptag_str, attrs=None, nsdecls=None):
     CT_Anchor, CT_Blip, CT_BlipFillProperties, CT_GraphicalObject,
     CT_GraphicalObjectData, CT_Inline, CT_NonVisualDrawingProps, CT_Picture,
     CT_PictureNonVisual, CT_Point2D, CT_PositiveSize2D, CT_ShapeProperties,
-    CT_Transform2D
+    CT_Transform2D, CT_WrapSquare, CT_WrapTopAndBottom
 )
-register_element_cls('a:blip',        CT_Blip)
-register_element_cls('a:ext',         CT_PositiveSize2D)
-register_element_cls('a:graphic',     CT_GraphicalObject)
-register_element_cls('a:graphicData', CT_GraphicalObjectData)
-register_element_cls('a:off',         CT_Point2D)
-register_element_cls('a:xfrm',        CT_Transform2D)
-register_element_cls('pic:blipFill',  CT_BlipFillProperties)
-register_element_cls('pic:cNvPr',     CT_NonVisualDrawingProps)
-register_element_cls('pic:nvPicPr',   CT_PictureNonVisual)
-register_element_cls('pic:pic',       CT_Picture)
-register_element_cls('pic:spPr',      CT_ShapeProperties)
-register_element_cls('wp:docPr',      CT_NonVisualDrawingProps)
-register_element_cls('wp:extent',     CT_PositiveSize2D)
-register_element_cls('wp:inline',     CT_Inline)
-register_element_cls('wp:anchor',     CT_Anchor)
+register_element_cls('a:blip',              CT_Blip)
+register_element_cls('a:ext',               CT_PositiveSize2D)
+register_element_cls('a:graphic',           CT_GraphicalObject)
+register_element_cls('a:graphicData',       CT_GraphicalObjectData)
+register_element_cls('a:off',               CT_Point2D)
+register_element_cls('a:xfrm',              CT_Transform2D)
+register_element_cls('pic:blipFill',        CT_BlipFillProperties)
+register_element_cls('pic:cNvPr',           CT_NonVisualDrawingProps)
+register_element_cls('pic:nvPicPr',         CT_PictureNonVisual)
+register_element_cls('pic:pic',             CT_Picture)
+register_element_cls('pic:spPr',            CT_ShapeProperties)
+register_element_cls('wp:docPr',            CT_NonVisualDrawingProps)
+register_element_cls('wp:extent',           CT_PositiveSize2D)
+register_element_cls('wp:inline',           CT_Inline)
+register_element_cls('wp:anchor',           CT_Anchor)
+register_element_cls('wp:wrapSquare',       CT_WrapSquare)
+register_element_cls('wp:wrapTopAndBottom', CT_WrapTopAndBottom)
 
 from .styles import CT_LatentStyles, CT_LsdException, CT_Style, CT_Styles
 register_element_cls('w:basedOn',        CT_String)

diff --git a/docx/oxml/shape.py b/docx/oxml/shape.py
@@ -8,12 +8,13 @@
 from .ns import nsdecls
 from .simpletypes import (
     ST_Coordinate, ST_DrawingElementId, ST_PositiveCoordinate,
-    ST_RelationshipId, XsdString, XsdToken
+    ST_RelationshipId, XsdString, XsdStringEnumeration, XsdToken
 )
 from .xmlchemy import (
     BaseOxmlElement, OneAndOnlyOne, OptionalAttribute, RequiredAttribute,
     ZeroOrOne
 )
+from ..enum.shape import WRAP_SHAPE_TYPE
 
 
 class CT_Blip(BaseOxmlElement):
@@ -34,6 +35,29 @@ class CT_BlipFillProperties(BaseOxmlElement):
     ))
 
 
+class ST_WrapText(XsdStringEnumeration):
+    """
+    Valid values for `wrapText/@val`.
+    """
+    BOTHSIDES = 'bothSides'
+
+    _members = (BOTHSIDES,)
+
+
+class CT_WrapSquare(BaseOxmlElement):
+    """
+    ``<wp:wrapSquare wrapText="bothSides" />`` element for wrapping text
+    around a shape
+    """
+    wrapText = RequiredAttribute('wrapText', ST_WrapText)
+
+
+class CT_WrapTopAndBottom(BaseOxmlElement):
+    """
+    ``<wp:wrapTopAndBottom />`` element for setting image on its own.
+    """
+
+
 class CT_GraphicalObject(BaseOxmlElement):
     """
     ``<a:graphic>`` element, container for a DrawingML object
@@ -121,6 +145,7 @@ class CT_Anchor(CT_Inline):
     positionV = OneAndOnlyOne('wp:positionV')
     effectExtent = OneAndOnlyOne('wp:effectExtent')
     wrapSquare = ZeroOrOne('wp:wrapSquare')
+    wrapTopAndBottom = ZeroOrOne('wp:wrapTopAndBottom')
 
     @classmethod
     def new(cls, cx, cy, shape_id, pic, position, margin=None, wrap=None):
@@ -147,8 +172,15 @@ def new(cls, cx, cy, shape_id, pic, position, margin=None, wrap=None):
             anchor.set('distB', u"%d" % margin.get('bottom', 0))
             anchor.set('distL', u"%d" % margin.get('left', 0))
 
-        if wrap is not None:
-            anchor.wrapSquare.set('wrapText', wrap)
+        wrap_el = None
+        if wrap == WRAP_SHAPE_TYPE.wrapTopAndBottom:
+            wrap_el = anchor.get_or_add_wrapTopAndBottom()
+        elif wrap == WRAP_SHAPE_TYPE.wrapSquareBothSides:
+            wrap_el = anchor.get_or_add_wrapSquare()
+            wrap_el.wrapText = ST_WrapText.BOTHSIDES
+
+        if wrap_el is not None:
+            anchor.insert_element_before(wrap_el, 'wp:effectExtent')
 
         return anchor
 
@@ -167,7 +199,6 @@ def _inline_xml(cls):
             '  </wp:positionV>\n'
             '  <wp:extent cx="914400" cy="914400"/>\n'
             '  <wp:effectExtent l="0" t="0" r="0" b="0" />\n'
-            '  <wp:wrapSquare wrapText="bothSides" />\n'
             '  <wp:docPr id="666" name="unnamed"/>\n'
             '  <wp:cNvGraphicFramePr>\n'
             '    <a:graphicFrameLocks %s noChangeAspect="1"/>\n'