add custom part class registration

Steve Canny · Steve Canny · commit 7cf579b516b3 · 2013-08-06T01:35:04.000-07:00
opc.package.PartFactory allows a custom part class to be registered
for a content type and uses that class when called with a matching
content type.
diff --git a/opc/package.py b/opc/package.py
@@ -167,7 +167,12 @@ class PartFactory(object):
     Provides a way for client code to specify a subclass of |Part| to be
     constructed by |Unmarshaller| based on its content type.
     """
+    part_type_for = {}
+
     def __new__(cls, partname, content_type, blob):
+        if content_type in PartFactory.part_type_for:
+            CustomPartClass = PartFactory.part_type_for[content_type]
+            return CustomPartClass(partname, content_type, blob)
         return Part(partname, content_type, blob)
 
 
diff --git a/tests/test_package.py b/tests/test_package.py
@@ -13,6 +13,7 @@
 
 from mock import call, Mock, patch, PropertyMock
 
+from opc.constants import CONTENT_TYPE as CT
 from opc.oxml import CT_Relationships
 from opc.package import (
     OpcPackage, Part, PartFactory, _Relationship, RelationshipCollection,
@@ -192,6 +193,18 @@ def it_constructs_a_part_instance(self, Part_):
         Part_.assert_called_once_with(partname, content_type, blob)
         assert part == Part_.return_value
 
+    def it_constructs_custom_part_type_for_registered_content_types(self):
+        # mockery ----------------------
+        CustomPartClass = Mock(name='CustomPartClass')
+        partname, blob = (Mock(name='partname'), Mock(name='blob'))
+        # exercise ---------------------
+        PartFactory.part_type_for[CT.WML_DOCUMENT_MAIN] = CustomPartClass
+        part = PartFactory(partname, CT.WML_DOCUMENT_MAIN, blob)
+        # verify -----------------------
+        CustomPartClass.assert_called_once_with(partname,
+                                                CT.WML_DOCUMENT_MAIN, blob)
+        assert part is CustomPartClass.return_value
+
 
 class Describe_Relationship(object):
 
