On-demand loading of PIL plugins

imagekit/models.py

@@ -10,7 +10,9 @@ from django.db.models.signals import post_save, post_delete
 from django.utils.encoding import force_unicode, smart_str
 
 from imagekit.lib import Image, ImageFile
-from imagekit.utils import img_to_fobj, get_spec_files, open_image
+from imagekit.utils import img_to_fobj, get_spec_files, open_image, \
+        format_to_extension, extension_to_format, UnknownFormatError, \
+        UnknownExtensionError
 from imagekit.processors import ProcessorPipeline
 
 # Modify image file buffer size.
@@ -95,24 +97,26 @@ class ImageSpec(_ImageSpecMixin):
             dispatch_uid='%s.delete' % uid)
 
 
-def get_registered_extensions():
-    Image.preinit()
-    return Image.EXTENSION
-
-
 def _get_suggested_extension(name, format):
-    if format:
-        # Try to look up an extension by the format.
-        extensions = [k for k, v in get_registered_extensions().iteritems() \
-                if v == format.upper()]
-    else:
-        extensions = []
     original_extension = os.path.splitext(name)[1]
-    if not extensions or original_extension.lower() in extensions:
-        # If the original extension matches the format, use it.
+    try:
+        suggested_extension = format_to_extension(format)
+    except UnknownFormatError:
         extension = original_extension
     else:
-        extension = extensions[0]
+        if suggested_extension.lower() == original_extension.lower():
+            extension = original_extension
+        else:
+            try:
+                original_format = extension_to_format(original_extension)
+            except UnknownExtensionError:
+                extension = suggested_extension
+            else:
+                # If the formats match, give precedence to the original extension.
+                if format.lower() == original_format.lower():
+                    extension = original_extension
+                else:
+                    extension = suggested_extension
     return extension
 
 
@@ -129,7 +133,10 @@ class _ImageSpecFileMixin(object):
                 # The extension is explicit, so assume they want the matching format.
                 extension = os.path.splitext(filename)[1].lower()
                 # Try to guess the format from the extension.
-                format = get_registered_extensions().get(extension)
+                try:
+                    format = extension_to_format(extension)
+                except UnknownExtensionError:
+                    pass
         format = format or img.format or original_format or 'JPEG'
 
         if format != 'JPEG':

imagekit/utils.py

@@ -50,3 +50,82 @@ def _wrap_copy(f):
             pass
         return img
     return copy
+
+
+class UnknownExtensionError(Exception):
+    pass
+
+
+class UnknownFormatError(Exception):
+    pass
+
+
+_pil_init = 0
+
+
+def _preinit_pil():
+    """Loads the standard PIL file format drivers. Returns True if ``preinit()``
+    was called (and there's a potential that more drivers were loaded) or False
+    if there is no possibility that new drivers were loaded.
+
+    """
+    global _pil_init
+    if _pil_init < 1:
+        Image.preinit()
+        _pil_init = 1
+        return True
+    return False
+
+
+def _init_pil():
+    """Loads all PIL file format drivers. Returns True if ``init()`` was called
+    (and there's a potential that more drivers were loaded) or False if there is
+    no possibility that new drivers were loaded.
+
+    """
+    global _pil_init
+    _preinit_pil()
+    if _pil_init < 2:
+        Image.init()
+        _pil_init = 2
+        return True
+    return False
+
+
+def _extension_to_format(extension):
+    return Image.EXTENSION.get(extension.lower())
+
+
+def _format_to_extension(format):
+    for k, v in Image.EXTENSION.iteritems():
+        if v == format.upper():
+            return k
+    return None
+
+
+def extension_to_format(extension):
+    """Returns the format that corresponds to the provided extension.
+
+    """
+    format = _extension_to_format(extension)
+    if not format and _preinit_pil():
+        format = _extension_to_format(extension)
+    if not format and _init_pil():
+        format = _extension_to_format(extension)
+    if not format:
+        raise UnknownExtensionError(extension)
+    return format
+
+
+def format_to_extension(format):
+    """Returns the first extension that matches the provided format.
+
+    """
+    extension = _format_to_extension(format)
+    if not extension and _preinit_pil():
+        extension = _format_to_extension(format)
+    if not extension and _init_pil():
+        extension = _extension_to_format(extension)
+    if not extension:
+        raise UnknownFormatError(format)
+    return extension