img: add Png._parse_chunk_offsets()

Steve Canny · Steve Canny · commit a2782f5fe823 · 2014-02-27T12:08:34.000-08:00
diff --git a/docx/image/png.py b/docx/image/png.py
@@ -42,7 +42,30 @@ def _parse_chunk_offsets(cls, stream):
         *stream*. The offsets for a chunk type that may appear more than once
         are returned as a list regardless of their actual number in *stream*.
         """
-        raise NotImplementedError
+        chunk_offsets = {}
+        for chunk_type, offset in cls._iter_chunk_offsets(stream):
+            # this would need to be more sophisticated if we needed any of
+            # the chunks that can appear multiple times
+            chunk_offsets[chunk_type] = offset
+        return chunk_offsets
+
+    @staticmethod
+    def _iter_chunk_offsets(stream):
+        """
+        Generate a (chunk_type, chunk_offset) 2-tuple for each of the chunks
+        in the PNG image stream. Iteration stops after the IEND chunk is
+        returned.
+        """
+        chunk_offset = 8
+        while True:
+            chunk_data_len = stream.read_long(chunk_offset)
+            chunk_type = stream.read_str(4, chunk_offset, 4)
+            data_offset = chunk_offset + 8
+            yield chunk_type, data_offset
+            if chunk_type == 'IEND':
+                break
+            # incr offset for chunk len long, chunk type, chunk data, and CRC
+            chunk_offset += (4 + 4 + chunk_data_len + 4)
 
     @classmethod
     def _parse_chunks(cls, stream, chunk_offsets):
diff --git a/tests/image/test_png.py b/tests/image/test_png.py
@@ -10,11 +10,11 @@
 
 from docx.compat import BytesIO
 from docx.image.exceptions import InvalidImageStreamError
-from docx.image.helpers import StreamReader
+from docx.image.helpers import BIG_ENDIAN, StreamReader
 from docx.image.png import Png
 
 from ..unitutil import (
-    initializer_mock, class_mock, instance_mock, method_mock
+    initializer_mock, class_mock, instance_mock, method_mock, test_file
 )
 
 
@@ -45,6 +45,12 @@ def it_raises_on_png_having_no_IHDR_chunk(self, no_IHDR_fixture):
         with pytest.raises(InvalidImageStreamError):
             Png._parse_png_headers(stream_)
 
+    def it_parses_chunk_offsets_to_help_chunk_parser(
+            self, chunk_offset_fixture):
+        stream, expected_chunk_offsets = chunk_offset_fixture
+        chunk_offsets = Png._parse_chunk_offsets(stream)
+        assert chunk_offsets == expected_chunk_offsets
+
     # fixtures -------------------------------------------------------
 
     @pytest.fixture
@@ -59,6 +65,22 @@ def attrs_(self, request):
     def blob_(self, request):
         return instance_mock(request, bytes)
 
+    @pytest.fixture(params=[
+        ('150-dpi.png', {
+            'IHDR': 16, 'pHYs': 41, 'iCCP': 62, 'cHRM': 2713, 'IDAT': 2757,
+            'IEND': 146888}),
+        ('300-dpi.png', {
+            'IHDR': 16, 'pHYs': 41, 'tEXt': 62, 'IDAT': 99, 'IEND': 39917}),
+    ])
+    def chunk_offset_fixture(self, request):
+        filename, expected_chunk_offsets = request.param
+        path = test_file(filename)
+        with open(path, 'rb') as f:
+            blob = f.read()
+        stream = BytesIO(blob)
+        stream_rdr = StreamReader(stream, BIG_ENDIAN)
+        return stream_rdr, expected_chunk_offsets
+
     @pytest.fixture
     def chunk_offsets_(self, request):
         return dict()
