Added PrivateImageField (#55)

README.rst

@@ -69,6 +69,28 @@ The ``PrivateFileField`` also accepts the following kwargs:
 * ``max_file_size``: maximum file size in bytes. (1MB is 1024 * 1024)
 * ``storage``: the storage object to use, defaults to ``private_storage.storage.private_storage``
 
+
+Images
+------
+
+You can also use ``PrivateImageField`` which only allows you to upload images:
+
+.. code-block:: python
+
+    from django.db import models
+    from private_storage.fields import PrivateImageField
+
+    class MyModel(models.Model):
+        title = models.CharField("Title", max_length=200)
+        width = models.PositiveSmallIntegerField(default=0)
+        height = models.PositiveSmallIntegerField(default=0)
+        image = PrivateFileField("Image", width_field='width', height_field='height')
+
+The ``PrivateImageField`` also accepts the following kwargs on top of ``PrivateFileField``:
+
+* ``width_field``: optional field for that stores the width of the image
+* ``height_field``: optional field for that stores the height of the image
+
 Other topics
 ============
 

private_storage/fields.py

@@ -6,17 +6,18 @@
 import os
 import posixpath
 import warnings
-import sys
 
 import django
+from django.core import checks
 from django.core.exceptions import ValidationError
 from django.core.files.uploadedfile import UploadedFile
 from django.db import models
+from django.db.models.fields.files import ImageFileDescriptor, ImageFieldFile
+from django.forms import ImageField
 from django.template.defaultfilters import filesizeformat
 from django.utils.encoding import force_str, force_text
 from django.utils.translation import ugettext_lazy as _
 
-
 from .storage import private_storage
 
 logger = logging.getLogger(__name__)
@@ -108,3 +109,115 @@ def _get_clean_filename(self, filename):
         # As of Django 1.10+, file names are no longer cleaned locally, but cleaned by the storage.
         # This compatibility function makes sure all Django versions generate a safe filename.
         return os.path.normpath(self.storage.get_valid_name(os.path.basename(filename)))
+
+
+class PrivateImageField(PrivateFileField):
+    attr_class = ImageFieldFile
+    descriptor_class = ImageFileDescriptor
+    description = _("Image")
+
+    def __init__(self, verbose_name=None, name=None, width_field=None, height_field=None, **kwargs):
+        self.width_field, self.height_field = width_field, height_field
+        super().__init__(verbose_name, name, **kwargs)
+
+    def check(self, **kwargs):
+        return [
+            *super().check(**kwargs),
+            *self._check_image_library_installed(),
+        ]
+
+    def _check_image_library_installed(self):
+        try:
+            from PIL import Image  # NOQA
+        except ImportError:
+            return [
+                checks.Error(
+                    'Cannot use ImageField because Pillow is not installed.',
+                    hint=('Get Pillow at https://pypi.org/project/Pillow/ '
+                          'or run command "python -m pip install Pillow".'),
+                    obj=self,
+                    id='fields.E210',
+                )
+            ]
+        else:
+            return []
+
+    def deconstruct(self):
+        name, path, args, kwargs = super().deconstruct()
+        if self.width_field:
+            kwargs['width_field'] = self.width_field
+        if self.height_field:
+            kwargs['height_field'] = self.height_field
+        return name, path, args, kwargs
+
+    def contribute_to_class(self, cls, name, **kwargs):
+        super().contribute_to_class(cls, name, **kwargs)
+        # Attach update_dimension_fields so that dimension fields declared
+        # after their corresponding image field don't stay cleared by
+        # Model.__init__, see bug #11196.
+        # Only run post-initialization dimension update on non-abstract models
+        if not cls._meta.abstract:
+            models.signals.post_init.connect(self.update_dimension_fields, sender=cls)
+
+    def update_dimension_fields(self, instance, force=False, *args, **kwargs):
+        """
+        Update field's width and height fields, if defined.
+
+        This method is hooked up to model's post_init signal to update
+        dimensions after instantiating a model instance.  However, dimensions
+        won't be updated if the dimensions fields are already populated.  This
+        avoids unnecessary recalculation when loading an object from the
+        database.
+
+        Dimensions can be forced to update with force=True, which is how
+        ImageFileDescriptor.__set__ calls this method.
+        """
+        # Nothing to update if the field doesn't have dimension fields or if
+        # the field is deferred.
+        has_dimension_fields = self.width_field or self.height_field
+        if not has_dimension_fields or self.attname not in instance.__dict__:
+            return
+
+        # getattr will call the ImageFileDescriptor's __get__ method, which
+        # coerces the assigned value into an instance of self.attr_class
+        # (ImageFieldFile in this case).
+        file = getattr(instance, self.attname)
+
+        # Nothing to update if we have no file and not being forced to update.
+        if not file and not force:
+            return
+
+        dimension_fields_filled = not(
+            (self.width_field and not getattr(instance, self.width_field)) or
+            (self.height_field and not getattr(instance, self.height_field))
+        )
+        # When both dimension fields have values, we are most likely loading
+        # data from the database or updating an image field that already had
+        # an image stored.  In the first case, we don't want to update the
+        # dimension fields because we are already getting their values from the
+        # database.  In the second case, we do want to update the dimensions
+        # fields and will skip this return because force will be True since we
+        # were called from ImageFileDescriptor.__set__.
+        if dimension_fields_filled and not force:
+            return
+
+        # file should be an instance of ImageFieldFile or should be None.
+        if file:
+            width = file.width
+            height = file.height
+        else:
+            # No file, so clear dimensions fields.
+            width = None
+            height = None
+
+        # Update the width and height fields.
+        if self.width_field:
+            setattr(instance, self.width_field, width)
+        if self.height_field:
+            setattr(instance, self.height_field, height)
+
+    def formfield(self, **kwargs):
+        return super().formfield(**{
+            'form_class': ImageField,
+            **kwargs,
+        })