img: add image stream recognizer

Steve Canny · Steve Canny · commit a78fb402b369 · 2014-02-27T12:08:33.000-08:00
Image streams are recognized by image_cls_that_can_parse(stream), which
returns the Image subclass that knows how to parse the image format
contained in *stream*.
diff --git a/docx/exceptions.py b/docx/exceptions.py
@@ -0,0 +1,7 @@
+# encoding: utf-8
+
+
+class UnrecognizedImageError(Exception):
+    """
+    The provided image stream could not be recognized.
+    """
diff --git a/docx/image/__init__.py b/docx/image/__init__.py
@@ -18,16 +18,34 @@
     import Image as PIL_Image
 
 from docx.compat import BytesIO, is_string
+from docx.exceptions import UnrecognizedImageError
+from docx.image.png import Png
 from docx.opc.constants import CONTENT_TYPE as CT
 from docx.shared import lazyproperty
 
 
+SIGNATURES = (
+    # class, offset, signature_bytes
+    (Png, 0, b'\x89PNG\x0D\x0A\x1A\x0A'),
+)
+
+
 def image_cls_that_can_parse(stream):
     """
     Return the |Image| subclass that can parse the headers of the image file
     contained in *stream*.
     """
-    raise NotImplementedError
+    def read_32(stream):
+        stream.seek(0)
+        return stream.read(32)
+
+    header = read_32(stream)
+    for cls, offset, signature_bytes in SIGNATURES:
+        end = offset + len(signature_bytes)
+        found_bytes = header[offset:end]
+        if found_bytes == signature_bytes:
+            return cls
+    raise UnrecognizedImageError
 
 
 class Image_OLD(object):
diff --git a/docx/image/png.py b/docx/image/png.py
@@ -0,0 +1,9 @@
+# encoding: utf-8
+
+from __future__ import absolute_import, division, print_function
+
+
+class Png(object):
+    """
+    Image header parser for PNG images
+    """
diff --git a/tests/image/test_image.py b/tests/image/test_image.py
@@ -9,8 +9,10 @@
 import pytest
 
 from docx.compat import BytesIO
-from docx.image import Image_OLD
+from docx.exceptions import UnrecognizedImageError
+from docx.image import image_cls_that_can_parse, Image_OLD
 from docx.image.image import Image
+from docx.image.png import Png
 from docx.opc.constants import CONTENT_TYPE as CT
 
 from ..unitutil import (
@@ -19,6 +21,33 @@
 )
 
 
+class Describe_image_cls_that_can_parse(object):
+
+    def it_can_recognize_an_image_stream(self, image_cls_lookup_fixture):
+        stream, expected_class = image_cls_lookup_fixture
+        ImageSubclass = image_cls_that_can_parse(stream)
+        assert ImageSubclass is expected_class
+
+    def it_raises_on_unrecognized_image_stream(self):
+        stream = BytesIO(b'foobar 666 not an image stream')
+        with pytest.raises(UnrecognizedImageError):
+            image_cls_that_can_parse(stream)
+
+    # fixtures -------------------------------------------------------
+
+    @pytest.fixture(params=[
+        ('python-icon.png', Png),
+    ])
+    def image_cls_lookup_fixture(self, request):
+        image_filename, expected_class = request.param
+        image_path = test_file(image_filename)
+        with open(image_path, 'rb') as f:
+            blob = f.read()
+        image_stream = BytesIO(blob)
+        image_stream.seek(666)
+        return image_stream, expected_class
+
+
 class DescribeImage(object):
 
     def it_can_construct_from_an_image_path(self, from_path_fixture):
