edan_f0/oscar/apps/image/utils.py

import os
import zlib
import tarfile
import zipfile
import tempfile
import shutil

from PIL import Image as PImage

from django.core.files import File
from django.core.exceptions import FieldError

from oscar.core.loading import import_module
import_module('image.exceptions', ['ImageImportError', 'IdenticalImageError', 'InvalidImageArchive'], locals())
import_module('product.models', ['Item'], locals())
import_module('image.models', ['Image'], locals())


class Importer(object):
    
    allowed_extensions = ['.jpeg','.jpg','.gif','.png']
    
    def __init__(self, logger, field):
        self.logger = logger
        self._field = field
        
    def handle(self, dirname):
        stats = {
            'num_processed': 0,
            'num_skipped': 0,
            'num_invalid': 0         
        }
        image_dir, filenames = self._get_image_files(dirname)
        if image_dir:
            for filename in filenames:
                try:
                    lookup_value = self._get_lookup_value_from_filename(filename)
                    self._process_image(image_dir, filename, lookup_value)
                    stats['num_processed'] += 1
                except Item.MultipleObjectsReturned:
                    self.logger.warning("Multiple products matching %s='%s', skipping" % (self._field, lookup_value))
                    stats['num_skipped'] += 1
                except Item.DoesNotExist:
                    self.logger.warning("No item matching %s='%s'" % (self._field, lookup_value))
                    stats['num_skipped'] += 1
                except IdenticalImageError:
                    self.logger.warning(" - Identical image already exists for %s='%s', skipping" % (self._field, lookup_value))
                    stats['num_skipped'] += 1
                except IOError, e:
                    raise ImageImportError('%s is not a valid image (%s)' % (filename, e))    
                    stats['num_invalid'] += 1
                except FieldError, e:
                    raise ImageImportError(e)
                    self._process_image(image_dir, filename)
            if image_dir != dirname:
                shutil.rmtree(image_dir)
        else:
            raise InvalidImageArchive('%s is not a valid image archive' % dirname)
        self.logger.info("Finished image import: %(num_processed)d imported, %(num_skipped)d skipped" % stats)
        
    def _get_image_files(self, dirname):
        filenames = []
        image_dir = self._extract_images(dirname)
        if image_dir:
            for filename in os.listdir(image_dir):
                ext = os.path.splitext(filename)[1]
                if os.path.isfile(os.path.join(image_dir, filename)) and ext in self.allowed_extensions:
                    filenames.append(filename)
        return image_dir, filenames

    def _extract_images(self, dirname):
        '''
        Returns path to directory containing images in dirname if successful.
        Returns empty string if dirname does not exist, or could not be opened.
        Assumes that if dirname is a directory, then it contains images.
        If dirname is an archive (tar/zip file) then the path returned is to a 
        temporary directory that should be deleted when no longer required.
        '''
        if os.path.isdir(dirname):
            return dirname        

        ext = os.path.splitext(dirname)[1]
        if ext in ['.gz', '.tar']:
            image_dir = tempfile.mkdtemp()
            try:
                tar_file = tarfile.open(dirname)
                tar_file.extractall(image_dir)
                tar_file.close()
                return image_dir
            except (tarfile.TarError, zlib.error):
                return ""
        elif ext == '.zip':
            image_dir = tempfile.mkdtemp()
            try:
                zip_file = zipfile.ZipFile(dirname)
                zip_file.extractall(image_dir) 
                zip_file.close()               
                return image_dir
            except (zlib.error, zipfile.BadZipfile, zipfile.LargeZipFile):
                return ""
        # unknown archive - perhaps this should be treated differently
        return ""
                
    def _process_image(self, dirname, filename, lookup_value):
        file_path = os.path.join(dirname, filename)
        trial_image = PImage.open(file_path)
        trial_image.verify()
        
        kwargs = {self._field: lookup_value}
        item = Item._default_manager.get(**kwargs)
        
        new_data = open(file_path).read()
        next_index = 0
        for existing in item.images.all():
            next_index = existing.display_order + 1
            try:
                if new_data == existing.original.read():
                    raise IdenticalImageError()
            except IOError:
                # File probably doesn't exist
                existing.delete()
            
        new_file = File(open(file_path))
        im = Image(product=item, display_order=next_index)
        im.original.save(filename, new_file, save=False)
        im.save()
        self.logger.info(' - Image added to "%s"' % item)
        
    def _fetch_item(self, filename):
        kwargs = {self._field: self._get_lookup_value_from_filename(filename)}
        return Item._default_manager.get(**kwargs)
    
    def _get_lookup_value_from_filename(self, filename):
        return os.path.splitext(filename)[0]