img: add _Chunks.IHDR and .pHYs

Steve Canny · Steve Canny · commit b196d4ba32d7 · 2014-03-01T01:35:30.000-08:00
diff --git a/docx/image/png.py b/docx/image/png.py
@@ -217,25 +217,47 @@ class _Chunks(object):
     """
     Collection of the chunks parsed from a PNG image stream
     """
-    def __init__(self, chunk_lst):
+    def __init__(self, chunk_iterable):
         super(_Chunks, self).__init__()
-        self._chunks = chunk_lst
+        self._chunks = list(chunk_iterable)
 
     @classmethod
     def from_stream(cls, stream):
         """
         Return a |_Chunks| instance containing the PNG chunks in *stream*.
         """
         chunk_parser = _ChunkParser.from_stream(stream)
-        chunk_lst = [chunk for chunk in chunk_parser.iter_chunks()]
-        return cls(chunk_lst)
+        chunks = [chunk for chunk in chunk_parser.iter_chunks()]
+        return cls(chunks)
 
     @property
     def IHDR(self):
         """
         IHDR chunk in PNG image
         """
-        raise NotImplementedError
+        match = lambda chunk: chunk.type_name == PNG_CHUNK_TYPE.IHDR
+        IHDR = self._find_first(match)
+        if IHDR is None:
+            raise InvalidImageStreamError('no IHDR chunk in PNG image')
+        return IHDR
+
+    @property
+    def pHYs(self):
+        """
+        pHYs chunk in PNG image, or |None| if not present
+        """
+        match = lambda chunk: chunk.type_name == PNG_CHUNK_TYPE.pHYs
+        return self._find_first(match)
+
+    def _find_first(self, match):
+        """
+        Return first chunk in stream order returning True for function
+        *match*.
+        """
+        for chunk in self._chunks:
+            if match(chunk):
+                return chunk
+        return None
 
 
 class _ChunkParser(object):
diff --git a/tests/image/test_png.py b/tests/image/test_png.py
@@ -316,6 +316,19 @@ def it_can_construct_from_a_stream(self, from_stream_fixture):
         _Chunks__init_.assert_called_once_with(chunk_lst)
         assert isinstance(chunks, _Chunks)
 
+    def it_provides_access_to_the_IHDR_chunk(self, IHDR_fixture):
+        chunks, IHDR_chunk_ = IHDR_fixture
+        assert chunks.IHDR == IHDR_chunk_
+
+    def it_provides_access_to_the_pHYs_chunk(self, pHYs_fixture):
+        chunks, expected_chunk = pHYs_fixture
+        assert chunks.pHYs == expected_chunk
+
+    def it_raises_if_theres_no_IHDR_chunk(self, no_IHDR_fixture):
+        chunks = no_IHDR_fixture
+        with pytest.raises(InvalidImageStreamError):
+            chunks.IHDR
+
     # fixtures -------------------------------------------------------
 
     @pytest.fixture
@@ -341,6 +354,40 @@ def chunk_parser_(self, request):
     def _Chunks__init_(self, request):
         return initializer_mock(request, _Chunks)
 
+    @pytest.fixture
+    def IHDR_fixture(self, IHDR_chunk_, pHYs_chunk_):
+        chunks = (IHDR_chunk_, pHYs_chunk_)
+        chunks = _Chunks(chunks)
+        return chunks, IHDR_chunk_
+
+    @pytest.fixture
+    def IHDR_chunk_(self, request):
+        return instance_mock(
+            request, _IHDRChunk, type_name=PNG_CHUNK_TYPE.IHDR
+        )
+
+    @pytest.fixture
+    def no_IHDR_fixture(self, pHYs_chunk_):
+        chunks = (pHYs_chunk_,)
+        chunks = _Chunks(chunks)
+        return chunks
+
+    @pytest.fixture
+    def pHYs_chunk_(self, request):
+        return instance_mock(
+            request, _pHYsChunk, type_name=PNG_CHUNK_TYPE.pHYs
+        )
+
+    @pytest.fixture(params=[True, False])
+    def pHYs_fixture(self, request, IHDR_chunk_, pHYs_chunk_):
+        has_pHYs_chunk = request.param
+        chunks = [IHDR_chunk_]
+        if has_pHYs_chunk:
+            chunks.append(pHYs_chunk_)
+        expected_chunk = pHYs_chunk_ if has_pHYs_chunk else None
+        chunks = _Chunks(chunks)
+        return chunks, expected_chunk
+
     @pytest.fixture
     def stream_(self, request):
         return instance_mock(request, BytesIO)
