img: add ImagePart part class selector

Steve Canny · Steve Canny · commit bb8e602d7cad · 2013-12-28T22:39:43.000-08:00
Added part_class_selector function that selects ImagePart class for
parts with reltype RT.IMAGE.
diff --git a/docx/__init__.py b/docx/__init__.py
@@ -7,11 +7,20 @@
 
 # register custom Part classes with opc package reader
 
-from docx.opc.constants import CONTENT_TYPE as CT
+from docx.opc.constants import CONTENT_TYPE as CT, RELATIONSHIP_TYPE as RT
 from docx.opc.package import PartFactory
 
 from docx.parts.document import DocumentPart
+from docx.parts.image import ImagePart
 
+
+def part_class_selector(content_type, reltype):
+    if reltype == RT.IMAGE:
+        return ImagePart
+    return None
+
+
+PartFactory.part_class_selector = part_class_selector
 PartFactory.part_type_for[CT.WML_DOCUMENT_MAIN] = DocumentPart
 
-del CT, DocumentPart, PartFactory
+del CT, DocumentPart, PartFactory, part_class_selector
diff --git a/docx/opc/package.py b/docx/opc/package.py
@@ -334,7 +334,8 @@ class PartFactory(object):
     def __new__(cls, partname, content_type, reltype, blob, package):
         PartClass = None
         if cls.part_class_selector is not None:
-            PartClass = cls.part_class_selector(content_type, reltype)
+            part_class_selector = cls.part_class_selector.__func__
+            PartClass = part_class_selector(content_type, reltype)
         if PartClass is None:
             PartClass = cls._part_cls_for(content_type)
         return PartClass.load(partname, content_type, blob, package)
diff --git a/tests/opc/test_package.py b/tests/opc/test_package.py
@@ -390,7 +390,9 @@ def it_constructs_part_from_selector_if_defined(
         PartFactory.part_class_selector = cls_selector_fn_
         part = PartFactory(partname, content_type, reltype, blob, package)
         # verify -----------------------
-        cls_selector_fn_.assert_called_once_with(content_type, reltype)
+        cls_selector_fn_.__func__.assert_called_once_with(
+            content_type, reltype
+        )
         CustomPartClass_.load.assert_called_once_with(
             partname, content_type, blob, package
         )
@@ -444,7 +446,11 @@ def reset_part_class_selector():
 
     @pytest.fixture
     def cls_selector_fn_(self, request, CustomPartClass_):
-        return loose_mock(request, return_value=CustomPartClass_)
+        cls_selector_fn_ = loose_mock(request)
+        cls_selector_fn_.__func__ = loose_mock(
+            request, name='__func__', return_value=CustomPartClass_
+        )
+        return cls_selector_fn_
 
     @pytest.fixture
     def content_type_(self, request):
diff --git a/tests/parts/test_document.py b/tests/parts/test_document.py
@@ -219,7 +219,7 @@ def oxml_fromstring_(self, request):
         return function_mock(request, 'docx.parts.document.oxml_fromstring')
 
     @pytest.fixture
-    def package_(self, request, image_parts_):
+    def package_(self, request):
         return instance_mock(request, Package)
 
     @pytest.fixture
diff --git a/tests/parts/test_image.py b/tests/parts/test_image.py
@@ -6,9 +6,15 @@
 
 from __future__ import absolute_import, print_function, unicode_literals
 
-from docx.parts.image import Image
+import pytest
 
-from ..unitutil import test_file
+from docx.opc.constants import CONTENT_TYPE as CT, RELATIONSHIP_TYPE as RT
+from docx.opc.package import PartFactory
+from docx.opc.packuri import PackURI
+from docx.package import Package
+from docx.parts.image import Image, ImagePart
+
+from ..unitutil import instance_mock, method_mock, test_file
 
 
 class DescribeImage(object):
@@ -27,3 +33,50 @@ def it_can_construct_from_an_image_stream(self):
         assert isinstance(image, Image)
         assert image.sha1 == '79769f1e202add2e963158b532e36c2c0f76a70c'
         assert image.filename is None
+
+
+class DescribeImagePart(object):
+
+    def it_is_used_by_PartFactory_to_construct_image_part(self, load_fixture):
+        # fixture ----------------------
+        image_part_load_, partname_, blob_, package_, image_part_ = (
+            load_fixture
+        )
+        content_type = CT.JPEG
+        reltype = RT.IMAGE
+        # exercise ---------------------
+        part = PartFactory(partname_, content_type, reltype, blob_, package_)
+        # verify -----------------------
+        image_part_load_.assert_called_once_with(
+            partname_, content_type, blob_, package_
+        )
+        assert part is image_part_
+
+    # fixtures -------------------------------------------------------
+
+    @pytest.fixture
+    def blob_(self, request):
+        return instance_mock(request, str)
+
+    @pytest.fixture
+    def image_part_(self, request):
+        return instance_mock(request, ImagePart)
+
+    @pytest.fixture
+    def image_part_load_(self, request, image_part_):
+        return method_mock(
+            request, ImagePart, 'load', return_value=image_part_
+        )
+
+    @pytest.fixture
+    def load_fixture(
+            self, image_part_load_, partname_, blob_, package_, image_part_):
+        return image_part_load_, partname_, blob_, package_, image_part_
+
+    @pytest.fixture
+    def package_(self, request):
+        return instance_mock(request, Package)
+
+    @pytest.fixture
+    def partname_(self, request):
+        return instance_mock(request, PackURI)
