shp: add InlineShapes.__iter__()

Steve Canny · Steve Canny · commit 85be65c8f8e6 · 2013-12-24T22:39:00.000-08:00
diff --git a/docx/parts.py b/docx/parts.py
@@ -104,6 +104,9 @@ class InlineShape(object):
     Proxy for an ``<wp:inline>`` element, representing the container for an
     inline graphical object.
     """
+    def __init__(self, inline):
+        super(InlineShape, self).__init__()
+        self._inline = inline
 
 
 class InlineShapes(object):
@@ -115,8 +118,14 @@ def __init__(self, body_elm):
         super(InlineShapes, self).__init__()
         self._body = body_elm
 
+    def __iter__(self):
+        return (InlineShape(inline) for inline in self._inline_lst)
+
     def __len__(self):
+        return len(self._inline_lst)
+
+    @property
+    def _inline_lst(self):
         body = self._body
         xpath = './w:p/w:r/w:drawing/wp:inline'
-        inline_elms = body.xpath(xpath, namespaces=nsmap)
-        return len(inline_elms)
+        return body.xpath(xpath, namespaces=nsmap)
diff --git a/features/steps/shape.py b/features/steps/shape.py
@@ -26,7 +26,7 @@ def given_a_document_containing_two_inline_shapes(context):
 def given_inline_shape_collection_containing_two_shapes(context):
     docx_path = test_docx('shp-inline-shape-access')
     document = Document(docx_path)
-    context.inline_shapes = document.body.inline_shapes
+    context.inline_shapes = document.inline_shapes
 
 
 # then =====================================================
diff --git a/tests/test_parts.py b/tests/test_parts.py
@@ -11,7 +11,7 @@
 from mock import Mock
 
 from docx.oxml.parts import CT_Document
-from docx.parts import _Body, _Document, InlineShapes
+from docx.parts import _Body, _Document, InlineShape, InlineShapes
 from docx.table import Table
 from docx.text import Paragraph
 
@@ -264,6 +264,15 @@ def it_knows_how_many_inline_shapes_it_contains(
         print(inline_shapes._body.xml)
         assert len(inline_shapes) == inline_shape_count
 
+    def it_can_iterate_over_its_InlineShape_instances(
+            self, inline_shapes_fixture):
+        inline_shapes, inline_shape_count = inline_shapes_fixture
+        actual_count = 0
+        for inline_shape in inline_shapes:
+            assert isinstance(inline_shape, InlineShape)
+            actual_count += 1
+        assert actual_count == inline_shape_count
+
     # fixtures -------------------------------------------------------
 
     @pytest.fixture
