style: add CT_Styles.add_style_of_type()

Steve Canny · Steve Canny · commit ade2ff4a2f16 · 2015-02-08T00:01:39.000-08:00
diff --git a/docx/oxml/styles.py b/docx/oxml/styles.py
@@ -11,6 +11,25 @@
 )
 
 
+def styleId_from_name(name):
+    """
+    Return the style id corresponding to *name*, taking into account
+    special-case names such as 'Heading 1'.
+    """
+    return {
+        'caption':   'Caption',
+        'heading 1': 'Heading1',
+        'heading 2': 'Heading2',
+        'heading 3': 'Heading3',
+        'heading 4': 'Heading4',
+        'heading 5': 'Heading5',
+        'heading 6': 'Heading6',
+        'heading 7': 'Heading7',
+        'heading 8': 'Heading8',
+        'heading 9': 'Heading9',
+    }.get(name, name.replace(' ', ''))
+
+
 class CT_Style(BaseOxmlElement):
     """
     A ``<w:style>`` element, representing a style definition
@@ -27,6 +46,7 @@ class CT_Style(BaseOxmlElement):
     type = OptionalAttribute('w:type', WD_STYLE_TYPE)
     styleId = OptionalAttribute('w:styleId', ST_String)
     default = OptionalAttribute('w:default', ST_OnOff)
+    customStyle = OptionalAttribute('w:customStyle', ST_OnOff)
     del _tag_seq
 
     @property
@@ -60,7 +80,12 @@ def add_style_of_type(self, name, style_type, builtin):
         *style_type*. `w:style/@customStyle` is set based on the value of
         *builtin*.
         """
-        raise NotImplementedError
+        style = self.add_style()
+        style.type = style_type
+        style.customStyle = None if builtin else True
+        style.styleId = styleId_from_name(name)
+        style.name_val = name
+        return style
 
     def default_for(self, style_type):
         """
diff --git a/tests/oxml/test_styles.py b/tests/oxml/test_styles.py
@@ -0,0 +1,40 @@
+# encoding: utf-8
+
+"""
+Test suite for the docx.oxml.styles module.
+"""
+
+from __future__ import (
+    absolute_import, division, print_function, unicode_literals
+)
+
+import pytest
+
+from docx.enum.style import WD_STYLE_TYPE
+
+from ..unitutil.cxml import element, xml
+
+
+class DescribeCT_Styles(object):
+
+    def it_can_add_a_style_of_type(self, add_fixture):
+        styles, name, style_type, builtin, expected_xml = add_fixture
+        style = styles.add_style_of_type(name, style_type, builtin)
+        assert styles.xml == expected_xml
+        assert style is styles[-1]
+
+    # fixtures -------------------------------------------------------
+
+    @pytest.fixture(params=[
+        ('w:styles', 'Foo Bar',   WD_STYLE_TYPE.LIST,      False,
+         'w:styles/w:style{w:type=numbering,w:customStyle=1,w:styleId=FooBar'
+         '}/w:name{w:val=Foo Bar}'),
+        ('w:styles', 'heading 1', WD_STYLE_TYPE.PARAGRAPH, True,
+         'w:styles/w:style{w:type=paragraph,w:styleId=Heading1}/w:name{w:val'
+         '=heading 1}'),
+    ])
+    def add_fixture(self, request):
+        styles_cxml, name, style_type, builtin, expected_cxml = request.param
+        styles = element(styles_cxml)
+        expected_xml = xml(expected_cxml)
+        return styles, name, style_type, builtin, expected_xml
