storage.py

# encoding=utf-8
"""
Author: Jökull Sólberg Auðunsson (solberg.is)

"""

from django.core.files.storage import get_storage_class
from django.db.models import ImageField
from uuid import uuid4

from PIL import ImageFile as PILImageFile
import os

BaseStorage = get_storage_class()

class ImageStorage(BaseStorage):
    """ Distributes files evenly to a filesystem. Tries to correct bad
    filename extensions.
    
    Example:
    
        >>> image = ImageField(upload_to='uploads/post/', storage=ImageStorage())
        
    Results in [IMG_0001.JPG -> ./uploads/post/69/37a5ff8c.jpg]
    
    """
    
    def find_extension(self, content):
        # Use PIL to determine filetype
        p = PILImageFile.Parser()
    
        for data in content.chunks():
            p.feed(data)
            if p.image:
                format = p.image.format
                break
        else:
            format = 'jpg' # Wild guess
            # TODO raise error?
        format = format.lower()
        if format == 'jpeg':
            format = 'jpg'
        try:
            p.close()
        except IOError:
            pass
        content.seek(0)
        return format
    
    def _save(self, name, content):
        extension = self.find_extension(content)
        basename, _ = os.path.splitext(name)
        name = '.'.join([basename, extension])
        return super(ImageStorage, self)._save(name, content)
    
    def get_valid_name(self, name):
        name = list(uuid4().hex[:10])
        name.insert(2, '/') # entropy[:2] + '/' + entropy[2:]
        return ''.join(name)