style: add Styles.add_style()

Steve Canny · Steve Canny · commit 03c2fdb575ab · 2015-02-08T00:01:39.000-08:00
diff --git a/docx/oxml/parts/styles.py b/docx/oxml/parts/styles.py
@@ -54,6 +54,14 @@ class CT_Styles(BaseOxmlElement):
     """
     style = ZeroOrMore('w:style', successors=())
 
+    def add_style_of_type(self, name, style_type, builtin):
+        """
+        Return a newly added `w:style` element having *name* and
+        *style_type*. `w:style/@customStyle` is set based on the value of
+        *builtin*.
+        """
+        raise NotImplementedError
+
     def default_for(self, style_type):
         """
         Return `w:style[@w:type="*{style_type}*][-1]` or |None| if not found.
diff --git a/docx/styles/styles.py b/docx/styles/styles.py
@@ -48,6 +48,18 @@ def __iter__(self):
     def __len__(self):
         return len(self._element.style_lst)
 
+    def add_style(self, name, style_type, builtin=False):
+        """
+        Return a newly added style object of *style_type* and identified
+        by *name*. A builtin style can be defined by passing True for the
+        optional *builtin* argument.
+        """
+        name = self._translate_special_case_names(name)
+        if name in self:
+            raise ValueError("document already contains style '%s'" % name)
+        style = self._element.add_style_of_type(name, style_type, builtin)
+        return StyleFactory(style)
+
     def default(self, style_type):
         """
         Return the default style for *style_type* or |None| if no default is
diff --git a/tests/styles/test_styles.py b/tests/styles/test_styles.py
@@ -11,6 +11,7 @@
 import pytest
 
 from docx.enum.style import WD_STYLE_TYPE
+from docx.oxml.parts.styles import CT_Style, CT_Styles
 from docx.styles.style import BaseStyle
 from docx.styles.styles import Styles
 
@@ -54,6 +55,23 @@ def it_raises_on_style_not_found(self, get_raises_fixture):
         with pytest.raises(KeyError):
             styles[key]
 
+    def it_can_add_a_new_style(self, add_fixture):
+        styles, name, style_type, builtin = add_fixture[:4]
+        name_, StyleFactory_, style_elm_, style_ = add_fixture[4:]
+
+        style = styles.add_style(name, style_type, builtin)
+
+        styles._element.add_style_of_type.assert_called_once_with(
+            name_, style_type, builtin
+        )
+        StyleFactory_.assert_called_once_with(style_elm_)
+        assert style is style_
+
+    def it_raises_when_style_name_already_used(self, add_raises_fixture):
+        styles, name = add_raises_fixture
+        with pytest.raises(ValueError):
+            styles.add_style(name, None)
+
     def it_can_get_the_default_style_for_a_type(self, default_fixture):
         styles, style_type, StyleFactory_ = default_fixture[:3]
         StyleFactory_calls, style_ = default_fixture[3:]
@@ -119,6 +137,28 @@ def it_raises_on_style_type_mismatch(self, id_style_raises_fixture):
 
     # fixture --------------------------------------------------------
 
+    @pytest.fixture(params=[
+        ('Foo Bar',   'Foo Bar',   WD_STYLE_TYPE.CHARACTER, False),
+        ('Heading 1', 'heading 1', WD_STYLE_TYPE.PARAGRAPH, True),
+    ])
+    def add_fixture(self, request, styles_elm_, _getitem_, style_elm_,
+                    StyleFactory_, style_):
+        name, name_, style_type, builtin = request.param
+        styles = Styles(styles_elm_)
+        _getitem_.return_value = None
+        styles_elm_.add_style_of_type.return_value = style_elm_
+        StyleFactory_.return_value = style_
+        return (
+            styles, name, style_type, builtin, name_, StyleFactory_,
+            style_elm_, style_
+        )
+
+    @pytest.fixture
+    def add_raises_fixture(self, _getitem_):
+        styles = Styles(element('w:styles/w:style/w:name{w:val=heading 1}'))
+        name = 'Heading 1'
+        return styles, name
+
     @pytest.fixture(params=[
         ('w:styles',
          False, WD_STYLE_TYPE.CHARACTER),
@@ -319,3 +359,11 @@ def style_(self, request):
     @pytest.fixture
     def StyleFactory_(self, request):
         return function_mock(request, 'docx.styles.styles.StyleFactory')
+
+    @pytest.fixture
+    def style_elm_(self, request):
+        return instance_mock(request, CT_Style)
+
+    @pytest.fixture
+    def styles_elm_(self, request):
+        return instance_mock(request, CT_Styles)
