Merge pull request #32 from scattm/D/0.0.4/I9

[R/0.0.4/I9] Add option for limit image file size in the configuration #9

config.conf.dist

@@ -5,3 +5,4 @@ REDISLITE_WORKER_PID = '<path-to-your>/workers.pid'
 
 IMAGEBUTLER_MAX_THUMBNAIL = 150, 150
 IMAGEBUTLER_API_IMAGES_LIMIT = 100
+IMAGEBUTLER_MAX_IMAGE_SIZE = 1M

imagebutler/models.py

@@ -1,6 +1,7 @@
 """Docstring for image_butler.models module."""
 
 import datetime
+from copy import deepcopy
 from flask_login import UserMixin
 from . import utils
 from werkzeug.datastructures import FileStorage
@@ -103,6 +104,9 @@ class ImageModel(CustomModelMixin, db.Model):
                                nullable=True)
     file_checksum = db.Column('fileChecksum', db.String(64),
                               nullable=False)
+    file_size = db.Column('fileSize', db.Integer,
+                          nullable=False,
+                          default=0)
     is_deleted = db.Column('isDeleted', db.Boolean,
                            nullable=False, default=False)
     user_id = db.Column('userId', db.Integer,
@@ -136,13 +140,11 @@ def __init__(self, file, user_id, file_description=None):
 
         # Image processing
         image = Image.open(file.stream)
-        image_sio = utils.BytesIO()
-        image.save(image_sio, format=image.format)
         # Set image values into the model
-        self.file_exif = utils.get_image_exif(image)
-        self.file_content = image_sio.getvalue()
-        self.file_checksum = utils.get_checksum(self.file_content)
-        image_sio.close()
+        image, self.file_content, self.file_checksum, \
+            self.file_size, self.file_exif = \
+            utils.process_uploaded_image(image,
+                                         config['IMAGEBUTLER_MAX_IMAGE_SIZE'])
 
         # Set thumbnail into the model
         self.file_thumbnail = self.gen_thumbnail(image)
@@ -158,10 +160,9 @@ def gen_thumbnail(self, image_instance=None):
         :return: BytesIO object
         """
 
-        if not image_instance:
-            temp_image = Image.open(utils.BytesIO(self.file_content))
-        else:
-            temp_image = image_instance.copy()
+        # Image.copy() does not work since it do not copy the image format.
+        temp_image = deepcopy(image_instance) if image_instance \
+            else self.image
 
         image_sio = utils.BytesIO()
         temp_image.thumbnail(
@@ -172,16 +173,24 @@ def gen_thumbnail(self, image_instance=None):
         temp_image.close()
         return image_sio.getvalue()
 
+    @property
+    def image(self):
+        """Return the image instance from the database or from argument."""
+
+        return Image.open(utils.BytesIO(self.file_content))
+
     @property
     def serving_object(self):
         """The model will not direct serving data to Flask but rather a
         medium class."""
+
         return ImageServingObject(self.file_mime,
                                   self.file_content,
                                   self.file_thumbnail)
 
     def __repr__(self):
         """Print the ImageModel instance."""
+
         return '<Image {file_name} - User Id {user_id}>'.format(
             file_name=self.file_name,
             user_id=self.user_id

imagebutler/utils.py

@@ -3,6 +3,8 @@
 import uuid
 import base64
 import pickle
+import piexif
+from PIL import Image
 from io import BytesIO
 from Crypto import Random
 from Crypto.Hash import SHA256
@@ -70,18 +72,82 @@ def get_image_exif(image):
     """
 
     try:
-        exif = image._getexif()  # There is not public method for this.
+        exif = image.info['exif']
 
-        if exif is not None:
-            exif_io = BytesIO()
-            pickle.dump(exif, exif_io, pickle.HIGHEST_PROTOCOL)
-            exif_data = exif_io.getvalue()
-            exif_io.close()
-            return exif_data
+        serialized_exif = \
+            pickle.dumps(piexif.load(exif), pickle.HIGHEST_PROTOCOL)
+        return serialized_exif
 
+    except KeyError:
         return None
-    except AttributeError:
-        return None
+
+
+def process_uploaded_image(image, max_size=0):
+    """
+    Process uploaded and return necessary information.
+
+    For resizing old image, we don't do it now but maybe in later releases.
+    :param image: Image object.
+    :type image: Image.Image
+    :param max_size: max size from the configuration
+    :type max_size: str
+    :return: a tuple of processed image information.
+    """
+
+    image_sio = BytesIO()
+    image_exif = get_image_exif(image)
+    if image_exif:
+        image_exif_deserialized = pickle.loads(image_exif)
+        image.save(image_sio, format=image.format,
+                   exif=piexif.dump(image_exif_deserialized)
+                   )
+    else:
+        image_exif_deserialized = None
+        image.save(image_sio, format=image.format)
+    image_size = len(image_sio.getvalue())
+    image_dimension_size = image.size
+
+    # TODO: Move this to a utils function
+    max_size = str(max_size)
+    matched_config = re.match(r'^(\d+)M$', max_size)
+    if matched_config:
+        max_size = int(matched_config.group(1)) * 1024 * 1024
+    else:
+        matched_config = re.match(r'^(\d+)K$', max_size)
+        if matched_config:
+            max_size = int(matched_config.group(1)) * 1024
+        else:
+            max_size = 0
+
+    if 0 < max_size < image_size:
+        resize_ratio = (float(max_size) / image_size) ** 0.5
+        new_dimension_size = (
+            int(image_dimension_size[0] * resize_ratio),
+            int(image_dimension_size[1] * resize_ratio)
+        )
+        print(max_size, image_size, resize_ratio)
+        image.thumbnail(new_dimension_size, Image.ANTIALIAS)
+        image_dimension_size = image.size
+        image_sio.close()
+        image_sio = BytesIO()
+        if image_exif:
+            image_exif_deserialized["0th"][piexif.ImageIFD.XResolution] = (
+                image_dimension_size[0], 1
+            )
+            image_exif_deserialized["0th"][piexif.ImageIFD.YResolution] = (
+                image_dimension_size[1], 1
+            )
+            image.save(image_sio, format=image.format,
+                       exif=piexif.dump(image_exif_deserialized))
+        else:
+            image.save(image_sio, format=image.format)
+        image_size = len(image_sio.getvalue())
+
+    image_file_content = image_sio.getvalue()
+    image_checksum = get_checksum(image_file_content)
+    image_sio.close()
+
+    return image, image_file_content, image_checksum, image_size, image_exif
 
 
 def user_identity_check(user, password):

requirements.txt

@@ -17,7 +17,9 @@ Mako==1.0.7
 MarkupSafe==1.0
 olefile==0.44
 pep8==1.7.1
+piexif==1.0.13
 Pillow==4.3.0
+pluggy==0.6.0
 progressbar2==3.34.3
 psutil==5.4.0
 py==1.5.2
@@ -37,4 +39,6 @@ redislite==3.2.311
 rq==0.9.1
 six==1.11.0
 SQLAlchemy==1.1.14
+tox==2.9.1
+virtualenv==15.1.0
 Werkzeug==0.12.2

setup.py

@@ -40,7 +40,8 @@
         'Flask-Login',
         'pycrypto',
         'Pillow',
-        'progressbar2'
+        'progressbar2',
+        'piexif'
     ],
     classifiers=[
         'Development Status :: 2 - Pre-Alpha',

tests/test_utils.py

@@ -2,7 +2,10 @@
 import pytest
 import uuid
 import time
+import pickle
+import piexif
 from Crypto import Random
+from PIL import Image
 from imagebutler import utils
 
 
@@ -140,6 +143,74 @@ def test_generate_password():
 def test_get_image_exif(sample_pil_jpeg_object_with_exif,
                         sample_pil_jpeg_object_no_exif,
                         sample_pil_png_object):
-    assert utils.get_image_exif(sample_pil_jpeg_object_with_exif)
-    assert utils.get_image_exif(sample_pil_jpeg_object_no_exif) is None
+    """Test get_image_exif function to return exif data or None."""
+
     assert utils.get_image_exif(sample_pil_png_object) is None
+    assert utils.get_image_exif(sample_pil_jpeg_object_no_exif) is None
+
+    exif = utils.get_image_exif(sample_pil_jpeg_object_with_exif)
+    assert exif
+    # Ensure that the function return right format
+    deserialize_exif = pickle.loads(exif)
+    assert deserialize_exif
+    assert piexif.dump(deserialize_exif)
+
+
+def test_process_uploaded_image_with_exif(
+        sample_pil_jpeg_object_with_exif
+):
+    """Test process_uploaded_image with image having exif."""
+
+    ret_1 = utils.process_uploaded_image(
+        sample_pil_jpeg_object_with_exif,
+        '1M'
+    )
+    assert isinstance(ret_1[0], Image.Image)
+    assert ret_1[1]
+    assert ret_1[2]
+    assert ret_1[3] < 1 * 1024 * 1024
+    assert ret_1[4]
+
+    ret_2 = utils.process_uploaded_image(
+        sample_pil_jpeg_object_with_exif,
+        '512K'
+    )
+    # Reserve place for EXIF
+    assert isinstance(ret_2[0], Image.Image)
+    assert ret_2[1]
+    assert ret_2[2]
+    assert ret_2[3] < 512 * 1024 * 1.05
+    assert ret_2[4]
+
+    ret_3 = utils.process_uploaded_image(
+        sample_pil_jpeg_object_with_exif,
+        '0'
+    )
+    assert isinstance(ret_3[0], Image.Image)
+    assert ret_3[1]
+    assert ret_3[2]
+    assert ret_3[4]
+
+
+def test_process_uploaded_image_without_exif(
+        sample_pil_jpeg_object_no_exif
+):
+    """Test process_uploaded_image with image not having exif."""
+
+    ret_1 = utils.process_uploaded_image(
+        sample_pil_jpeg_object_no_exif,
+        '1M'
+    )
+    assert isinstance(ret_1[0], Image.Image)
+    assert ret_1[1]
+    assert ret_1[2]
+    assert ret_1[3] < 1 * 1024 * 1024
+    assert ret_1[4] is None
+
+    ret_2 = utils.process_uploaded_image(
+        sample_pil_jpeg_object_no_exif,
+        '256K'
+    )
+    # Reserve place for EXIF
+    assert ret_2[3] < 256 * 1024 * 1.05
+    assert ret_2[4] is None