Refactor image loading.

weasy/document.py

@@ -37,6 +37,7 @@ from .formatting_structure.build import build_formatting_structure
 from .layout import layout
 from . import draw
 from . import utils
+from . import images
 
 
 def make_parser():
@@ -186,15 +187,11 @@ class Document(object):
         return self._pages
 
     def get_image_surface_from_uri(self, uri):
-        if uri in self._image_cache:
-            return self._image_cache[uri]
-        try:
-            surface = draw.get_image_surface_from_uri(uri)
-        # TODO: have a more specific list of exception for network errors
-        # and PNG parsing errors.
-        except Exception:
-            surface = None
-        self._image_cache[uri] = surface
+        missing = object()
+        surface = self._image_cache.get(uri, missing)
+        if surface is missing:
+            surface = images.get_image_surface_from_uri(uri)
+            self._image_cache[uri] = surface
         return surface
 
 

weasy/draw.py

@@ -24,35 +24,11 @@ Module drawing documents.
 
 from __future__ import division
 import contextlib
-import urllib
 
 import cairo
-from StringIO import StringIO
 
 from .formatting_structure import boxes
 from .css.values import get_percentage_value
-from .utils import urlopen
-
-
-SUPPORTED_IMAGES = ['image/png', 'image/gif', 'image/jpeg', 'image/bmp']
-
-
-def get_image_surface_from_uri(uri):
-    """Get a :class:`cairo.ImageSurface`` from an image URI."""
-    file_like, mime_type, _charset = urlopen(uri)
-    # TODO: implement image type sniffing?
-    # http://www.w3.org/TR/html5/fetching-resources.html#content-type-sniffing:-image
-    if mime_type in SUPPORTED_IMAGES:
-        if mime_type == "image/png":
-            image = file_like
-        else:
-            from PIL import Image
-            pil_image = Image.open(StringIO(file_like.read()))
-            image = StringIO()
-            pil_image = pil_image.convert('RGBA')
-            pil_image.save(image, "PNG")
-            image.seek(0)
-        return cairo.ImageSurface.create_from_png(image)
 
 
 class CairoContext(cairo.Context):

weasy/images.py

@@ -0,0 +1,92 @@
+# coding: utf8
+
+#  WeasyPrint converts web documents (HTML, CSS, ...) to PDF.
+#  Copyright (C) 2011  Simon Sapin
+#
+#  This program is free software: you can redistribute it and/or modify
+#  it under the terms of the GNU Affero General Public License as
+#  published by the Free Software Foundation, either version 3 of the
+#  License, or (at your option) any later version.
+#
+#  This program is distributed in the hope that it will be useful,
+#  but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#  GNU Affero General Public License for more details.
+#
+#  You should have received a copy of the GNU Affero General Public License
+#  along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+"""
+Handle various image formats.
+"""
+
+from __future__ import division
+from StringIO import StringIO
+import urllib
+
+import cairo
+
+from .utils import urlopen
+
+
+SUPPORTED_FORMATS = {}
+
+
+def register_format(mime_type):
+    """Register a handler for a give image MIME type."""
+    def decorator(function):
+        SUPPORTED_FORMATS[mime_type] = function
+        return function
+    return decorator
+
+
+@register_format('image/png')
+def png_handler(file_like):
+    """Return a cairo Surface from a PNG byte stream."""
+    return cairo.ImageSurface.create_from_png(file_like)
+
+
+def fallback_handler(file_like):
+    """
+    Parse a byte stream with PIL and return a cairo Surface.
+
+    PIL supports many raster image formats and does not take a `format`
+    parameter, it guesses the format from the content.
+    """
+    try:
+        from PIL import Image
+    except ImportError:
+        try:
+            # It is sometimes installed with another name...
+            import Image
+        except ImportError:
+            return None  # PIL is not installed
+    if not (hasattr(file_like, 'seek') and hasattr(file_like, 'tell')):
+        # PIL likes to have these methods
+        file_like = StringIO(file_like.read())
+    png = StringIO()
+    image = Image.open(file_like)
+    image = image.convert('RGBA')
+    image.save(png, "PNG")
+    png.seek(0)
+    return png_handler(png)
+
+
+def get_image_surface_from_uri(uri):
+    """Get a :class:`cairo.Surface`` from an image URI."""
+    try:
+        file_like, mime_type, _charset = urlopen(uri)
+    except IOError:
+        # TODO: warn
+        return None
+    # TODO: implement image type sniffing?
+# http://www.w3.org/TR/html5/fetching-resources.html#content-type-sniffing:-image
+    handler = SUPPORTED_FORMATS.get(mime_type, fallback_handler)
+    try:
+        return handler(file_like)
+    except (IOError, MemoryError):
+        # Network or parsing error
+        # TODO: warn
+        return None
+    finally:
+        file_like.close()

weasy/tests/test_layout.py

@@ -736,8 +736,25 @@ def test_whitespace_processing():
 @SUITE.test
 def test_with_images():
     """Test that width, height and ratio of images are respected."""
-    # pattern.png is 4x4 px. Layout rules try to preserve the ratio, so
-    # the height should be 40px too:
+    page, = parse('<img src="pattern.png">')
+    html = page.root_box
+    body, = html.children
+    line, = body.children
+    img, = line.children
+    # pattern.png is 4x4 px.
+    assert img.width == 4
+    assert img.height == 4
+
+    # Try the same image in JPEG
+    page, = parse('<img src="pattern.jpg">')
+    html = page.root_box
+    body, = html.children
+    line, = body.children
+    img, = line.children
+    assert img.width == 4
+    assert img.height == 4
+
+    # Layout rules try to preserve the ratio, so the height should be 40px too:
     page, = parse('<img src="pattern.png" style="width: 40px">')
     html = page.root_box
     body, = html.children
@@ -745,6 +762,8 @@ def test_with_images():
     img, = line.children
     assert body.height == 40
     assert img.position_y == 0
+    assert img.width == 40
+    assert img.height == 40
 
     page, = parse('''
         <img src="pattern.png" style="width: 40px">

weasy/utils.py

@@ -59,6 +59,8 @@ class URLopener(urllib.FancyURLopener):
 
 def urlopen(url):
     """Fetch an URL and return ``(file_like, mime_type, charset)``.
+
+    It is the caller’s responsability to call ``file_like.close()``.
     """
     file_like = URLopener().open(url)
     info = file_like.info()
@@ -74,7 +76,7 @@ def urlopen(url):
     else:
         # Python 2
         charset = info.getparam('charset')
-    return file_like, mime_type, charset
+    return file_like.fp, mime_type, charset
 
 
 def urllib_fetcher(url):
@@ -88,4 +90,6 @@ def urllib_fetcher(url):
     if mime_type != 'text/css':
         # TODO: add a warning
         return None
-    return charset, file_like.read()
+    content = file_like.read()
+    file_like.close()
+    return charset, content