doc: add _Document.get_or_add_image_part()

Steve Canny · Steve Canny · commit 13e928f06371 · 2013-12-27T01:37:12.000-08:00
Along the way:

* introduce docx.package module and Package class
* trying out idea of suffixing part classes with Part, for clarity
diff --git a/docx/api.py b/docx/api.py
@@ -8,8 +8,8 @@
 
 import os
 
-from docx.opc.package import OpcPackage
 from docx.opc.constants import CONTENT_TYPE as CT
+from docx.package import Package
 
 
 thisdir = os.path.split(__file__)[0]
@@ -25,7 +25,7 @@ def Document(docx=None):
     """
     if docx is None:
         docx = _default_docx_path
-    pkg = OpcPackage.open(docx)
+    pkg = Package.open(docx)
     document_part = pkg.main_document
     if document_part.content_type != CT.WML_DOCUMENT_MAIN:
         tmpl = "file '%s' is not a Word file, content type is '%s'"
diff --git a/docx/package.py b/docx/package.py
@@ -0,0 +1,33 @@
+# encoding: utf-8
+
+"""
+WordprocessingML Package class and related objects
+"""
+
+from __future__ import absolute_import, print_function, unicode_literals
+
+from docx.opc.package import OpcPackage
+
+
+class Package(OpcPackage):
+    """
+    Customizations specific to a WordprocessingML package.
+    """
+    @property
+    def image_parts(self):
+        """
+        Collection of all image parts in this package.
+        """
+
+
+class ImageParts(object):
+    """
+    Collection of |ImagePart| instances containing all the image parts in the
+    package.
+    """
+    def get_or_add_image_part(self, image_descriptor):
+        """
+        Return an |ImagePart| instance containing the image identified by
+        *image_descriptor*, newly created if a matching one is not present in
+        the collection.
+        """
diff --git a/docx/parts/document.py b/docx/parts/document.py
@@ -5,6 +5,7 @@
 """
 
 from docx.enum.shape import WD_INLINE_SHAPE
+from docx.opc.constants import RELATIONSHIP_TYPE as RT
 from docx.opc.oxml import serialize_part_xml
 from docx.opc.package import Part
 from docx.oxml.shared import nsmap, oxml_fromstring
@@ -23,15 +24,18 @@ def __init__(self, partname, content_type, document_elm, package):
         )
         self._element = document_elm
 
-    def add_image(self, image_descriptor):
+    def get_or_add_image_part(self, image_descriptor):
         """
         Return an ``(image_part, rId)`` 2-tuple for the image identified by
         *image_descriptor*. *image_part* is an |Image| instance corresponding
-        to the image, newly created if not already present in document. *rId*
+        to the image, newly created if no matching image part is found. *rId*
         is the key for the relationship between this document part and the
         image part, reused if already present, newly created if not.
         """
-        raise NotImplementedError
+        image_parts = self._package.image_parts
+        image_part = image_parts.get_or_add_image_part(image_descriptor)
+        rId = self.relate_to(image_part, RT.IMAGE)
+        return (image_part, rId)
 
     @property
     def blob(self):
@@ -202,7 +206,7 @@ def add_picture(self, image_descriptor):
         end of the document. *image_descriptor* can be a path (a string) or a
         file-like object containing a binary image.
         """
-        rId, image = self.part.add_image(image_descriptor)
+        rId, image = self.part.get_or_add_image_part(image_descriptor)
         shape_id = self.part.next_id
         r = self._body.add_p().add_r()
         return InlineShape.new_picture(r, image, rId, shape_id)
diff --git a/docx/parts/image.py b/docx/parts/image.py
@@ -9,7 +9,7 @@
 from docx.opc.package import Part
 
 
-class Image(Part):
+class ImagePart(Part):
     """
     An image part. Corresponds to the target part of a relationship with type
     RELATIONSHIP_TYPE.IMAGE.
diff --git a/tests/parts/test_document.py b/tests/parts/test_document.py
@@ -11,11 +11,13 @@
 from mock import Mock
 
 from docx.enum.shape import WD_INLINE_SHAPE
+from docx.opc.constants import RELATIONSHIP_TYPE as RT
 from docx.oxml.parts.document import CT_Body, CT_Document
 from docx.oxml.shared import nsmap
 from docx.oxml.text import CT_R
+from docx.package import ImageParts, Package
 from docx.parts.document import _Body, _Document, InlineShape, InlineShapes
-from docx.parts.image import Image
+from docx.parts.image import ImagePart
 from docx.table import Table
 from docx.text import Paragraph
 
@@ -29,7 +31,7 @@
 from ..oxml.unitdata.text import a_p, a_sectPr, an_r
 from ..unitutil import (
     function_mock, class_mock, initializer_mock, instance_mock, loose_mock,
-    property_mock
+    method_mock, property_mock
 )
 
 
@@ -53,6 +55,18 @@ def it_can_be_constructed_by_opc_part_factory(
         )
         assert isinstance(doc, _Document)
 
+    def it_can_add_an_image_part_to_the_document(
+            self, get_or_add_image_fixture):
+        (document, image_descriptor_, image_parts_, relate_to_, image_part_,
+         rId_) = get_or_add_image_fixture
+        image_part, rId = document.get_or_add_image_part(image_descriptor_)
+        image_parts_.get_or_add_image_part.assert_called_once_with(
+            image_descriptor_
+        )
+        relate_to_.assert_called_once_with(image_part_, RT.IMAGE)
+        assert image_part is image_part_
+        assert rId == rId_
+
     def it_has_a_body(self, document_body_fixture):
         document, _Body_, body_elm = document_body_fixture
         _body = document.body
@@ -101,6 +115,30 @@ def document_body_fixture(self, request, _Body_):
     def _Body_(self, request):
         return class_mock(request, 'docx.parts.document._Body')
 
+    @pytest.fixture
+    def get_or_add_image_fixture(
+            self, request, package_, image_descriptor_, image_parts_,
+            relate_to_, image_part_, rId_):
+        document = _Document(None, None, None, package_)
+        return (
+            document, image_descriptor_, image_parts_, relate_to_,
+            image_part_, rId_
+        )
+
+    @pytest.fixture
+    def image_descriptor_(self, request):
+        return instance_mock(request, str)
+
+    @pytest.fixture
+    def image_part_(self, request):
+        return instance_mock(request, ImagePart)
+
+    @pytest.fixture
+    def image_parts_(self, request, image_part_):
+        image_parts_ = instance_mock(request, ImageParts)
+        image_parts_.get_or_add_image_part.return_value = image_part_
+        return image_parts_
+
     @pytest.fixture
     def init(self, request):
         return initializer_mock(request, _Document)
@@ -123,6 +161,22 @@ def inline_shapes_fixture(self, request, InlineShapes_):
     def oxml_fromstring_(self, request):
         return function_mock(request, 'docx.parts.document.oxml_fromstring')
 
+    @pytest.fixture
+    def package_(self, request, image_parts_):
+        package_ = instance_mock(request, Package)
+        package_.image_parts = image_parts_
+        return package_
+
+    @pytest.fixture
+    def relate_to_(self, request, rId_):
+        relate_to_ = method_mock(request, _Document, 'relate_to')
+        relate_to_.return_value = rId_
+        return relate_to_
+
+    @pytest.fixture
+    def rId_(self, request):
+        return instance_mock(request, str)
+
     @pytest.fixture
     def serialize_part_xml_(self, request):
         return function_mock(
@@ -374,10 +428,15 @@ def it_raises_on_indexed_access_out_of_range(
 
     def it_can_add_an_inline_picture_to_the_document(
             self, add_picture_fixture):
+        # fixture ----------------------
         (inline_shapes, image_descriptor_, document_, InlineShape_, r_,
          image_, rId_, shape_id_, new_picture_shape_) = add_picture_fixture
+        # exercise ---------------------
         picture_shape = inline_shapes.add_picture(image_descriptor_)
-        document_.add_image.assert_called_once_with(image_descriptor_)
+        # verify -----------------------
+        document_.get_or_add_image_part.assert_called_once_with(
+            image_descriptor_
+        )
         InlineShape_.new_picture.assert_called_once_with(
             r_, image_, rId_, shape_id_
         )
@@ -393,12 +452,12 @@ def it_knows_the_part_it_belongs_to(self, inline_shapes_with_parent_):
     @pytest.fixture
     def add_picture_fixture(
             self, request, body_, document_, image_descriptor_, InlineShape_,
-            r_, image_, rId_, shape_id_, new_picture_shape_):
+            r_, image_part_, rId_, shape_id_, new_picture_shape_):
         inline_shapes = InlineShapes(body_, None)
         property_mock(request, InlineShapes, 'part', return_value=document_)
         return (
             inline_shapes, image_descriptor_, document_, InlineShape_, r_,
-            image_, rId_, shape_id_, new_picture_shape_
+            image_part_, rId_, shape_id_, new_picture_shape_
         )
 
     @pytest.fixture
@@ -408,15 +467,15 @@ def body_(self, request, r_):
         return body_
 
     @pytest.fixture
-    def document_(self, request, rId_, image_, shape_id_):
+    def document_(self, request, rId_, image_part_, shape_id_):
         document_ = instance_mock(request, _Document, name='document_')
-        document_.add_image.return_value = rId_, image_
+        document_.get_or_add_image_part.return_value = rId_, image_part_
         document_.next_id = shape_id_
         return document_
 
     @pytest.fixture
-    def image_(self, request):
-        return instance_mock(request, Image)
+    def image_part_(self, request):
+        return instance_mock(request, ImagePart)
 
     @pytest.fixture
     def image_descriptor_(self, request):
diff --git a/tests/test_api.py b/tests/test_api.py
@@ -8,7 +8,7 @@
 
 from docx.api import Document, _Document
 from docx.opc.constants import CONTENT_TYPE as CT
-from docx.opc.package import OpcPackage
+from docx.package import Package
 from docx.parts.document import _Document as parts_Document, InlineShapes
 
 from .unitutil import class_mock, instance_mock, var_mock
@@ -17,18 +17,17 @@
 class DescribeDocument(object):
 
     def it_opens_a_docx_file(self, open_fixture):
-        docx_, OpcPackage_, _Document_, package, document_part = open_fixture
+        docx_, Package_, _Document_, package, document_part = open_fixture
         _document = Document(docx_)
-        OpcPackage_.open.assert_called_once_with(docx_)
+        Package_.open.assert_called_once_with(docx_)
         _Document_.assert_called_once_with(package, document_part)
         assert _document is _Document_.return_value
 
-    def it_uses_default_if_no_file_provided(self, OpcPackage_, default_docx_):
+    def it_uses_default_if_no_file_provided(self, Package_, default_docx_):
         Document()
-        OpcPackage_.open.assert_called_once_with(default_docx_)
+        Package_.open.assert_called_once_with(default_docx_)
 
-    def it_should_raise_if_not_a_Word_file(
-            self, OpcPackage_, package_, docx_):
+    def it_should_raise_if_not_a_Word_file(self, Package_, package_, docx_):
         package_.main_document.content_type = 'foobar'
         with pytest.raises(ValueError):
             Document(docx_)
@@ -53,21 +52,21 @@ def document_part_(self, request):
     def docx_(self, request):
         return instance_mock(request, str)
 
-    @pytest.fixture
-    def OpcPackage_(self, request, package_):
-        OpcPackage_ = class_mock(request, 'docx.api.OpcPackage')
-        OpcPackage_.open.return_value = package_
-        return OpcPackage_
-
     @pytest.fixture
     def open_fixture(
-            self, request, docx_, OpcPackage_, _Document_, package_,
+            self, request, docx_, Package_, _Document_, package_,
             document_part_):
-        return docx_, OpcPackage_, _Document_, package_, document_part_
+        return docx_, Package_, _Document_, package_, document_part_
+
+    @pytest.fixture
+    def Package_(self, request, package_):
+        Package_ = class_mock(request, 'docx.api.Package')
+        Package_.open.return_value = package_
+        return Package_
 
     @pytest.fixture
     def package_(self, request, document_part_):
-        package_ = instance_mock(request, OpcPackage)
+        package_ = instance_mock(request, Package)
         package_.main_document = document_part_
         return package_
 
@@ -118,7 +117,7 @@ def document_part_(self, request):
 
     @pytest.fixture
     def package_(self, request):
-        return instance_mock(request, OpcPackage)
+        return instance_mock(request, Package)
 
     @pytest.fixture
     def save_fixture(self, request, package_):
