boards/models/attachment/__init__.py

from itertools import zip_longest

import boards
from boards.models import STATUS_ARCHIVE
from django.core.files.images import get_image_dimensions
from django.db import models

from boards import utils
from boards.models.attachment.viewers import get_viewers, AbstractViewer, \
    FILE_TYPES_IMAGE
from boards.utils import get_upload_filename, get_extension, cached_result, \
        get_file_mimetype


class AttachmentManager(models.Manager):
    def create_with_hash(self, file):
        file_hash = utils.get_file_hash(file)
        attachment = self.get_existing_duplicate(file_hash, file)
        if not attachment:
            file_type = get_file_mimetype(file)
            attachment = self.create(file=file, mimetype=file_type,
                                     hash=file_hash)

        return attachment

    def create_from_url(self, url):
        existing = self.filter(url=url)
        if len(existing) > 0:
            attachment = existing[0]
        else:
            attachment = self.create(url=url)
        return attachment

    def get_random_images(self, count, tags=None):
        images = self.filter(mimetype__in=FILE_TYPES_IMAGE).exclude(
            attachment_posts__thread__status=STATUS_ARCHIVE)
        if tags is not None:
            images = images.filter(attachment_posts__threads__tags__in=tags)
        return images.order_by('?')[:count]

    def get_existing_duplicate(self, file_hash, file):
        """
        Gets an attachment with the same file if one exists.
        """
        existing = self.filter(hash=file_hash)
        attachment = None
        for existing_attachment in existing:
            existing_file = existing_attachment.file

            file_chunks = file.chunks()
            existing_file_chunks = existing_file.chunks()

            if self._compare_chunks(file_chunks, existing_file_chunks):
                attachment = existing_attachment
        return attachment

    def get_by_alias(self, name):
        pack_name, sticker_name = name.split('/')
        try:
            return AttachmentSticker.objects.get(name=sticker_name, stickerpack__name=pack_name).attachment
        except AttachmentSticker.DoesNotExist:
            return None

    def _compare_chunks(self, chunks1, chunks2):
        """
        Compares 2 chunks of different sizes (e.g. first chunk array contains
        all data in 1 chunk, and other one -- in a multiple of smaller ones.
        """
        equal = True

        position1 = 0
        position2 = 0
        chunk1 = None
        chunk2 = None
        chunk1ended = False
        chunk2ended = False
        while True:
            if not chunk1 or len(chunk1) <= position1:
                try:
                    chunk1 = chunks1.__next__()
                    position1 = 0
                except StopIteration:
                    chunk1ended = True
            if not chunk2 or len(chunk2) <= position2:
                try:
                    chunk2 = chunks2.__next__()
                    position2 = 0
                except StopIteration:
                    chunk2ended = True

            if chunk1ended and chunk2ended:
                # Same size chunksm checked for equality previously
                break
            elif chunk1ended or chunk2ended:
                # Different size chunks, not equal
                equal = False
                break
            elif chunk1[position1] != chunk2[position2]:
                # Different bytes, not equal
                equal = False
                break
            else:
                position1 += 1
                position2 += 1
        return equal


class Attachment(models.Model):
    objects = AttachmentManager()

    class Meta:
        app_label = 'boards'
        ordering = ('id',)

    file = models.FileField(upload_to=get_upload_filename, null=True)
    mimetype = models.CharField(max_length=200, null=True)
    hash = models.CharField(max_length=36, null=True)
    url = models.TextField(blank=True, default='')

    def get_view(self):
        file_viewer = None
        for viewer in get_viewers():
            if viewer.supports(self.mimetype):
                file_viewer = viewer
                break
        if file_viewer is None:
            file_viewer = AbstractViewer

        return file_viewer(self.file, self.mimetype, self.id, self.url).get_view()

    def __str__(self):
        return self.url or self.file.url

    def get_random_associated_post(self):
        posts = boards.models.Post.objects.filter(attachments__in=[self])
        return posts.order_by('?').first()

    @cached_result()
    def get_size(self):
        if self.file:
            if self.mimetype in FILE_TYPES_IMAGE:
                return get_image_dimensions(self.file)
            else:
                return 200, 150

    def get_thumb_url(self):
        split = self.file.url.rsplit('.', 1)
        w, h = 200, 150
        return '%s.%sx%s.%s' % (split[0], w, h, split[1])

    @cached_result()
    def get_preview_size(self):
        size = 200, 150
        if self.mimetype in FILE_TYPES_IMAGE:
            preview_path = self.file.path.replace('.', '.200x150.')
            try:
                size = get_image_dimensions(preview_path)
            except Exception:
                pass

        return size

    def is_internal(self):
        return self.url is None or len(self.url) == 0


class StickerPack(models.Model):
    name = models.TextField(unique=True)
    tripcode = models.TextField(blank=True)

    def __str__(self):
        return self.name


class AttachmentSticker(models.Model):
    attachment = models.ForeignKey('Attachment', on_delete=models.CASCADE)
    name = models.TextField(unique=True)
    stickerpack = models.ForeignKey('StickerPack', on_delete=models.CASCADE)

    def __str__(self):
        # Local stickers do not have a sticker pack
        if hasattr(self, 'stickerpack'):
            return '{}/{}'.format(str(self.stickerpack), self.name)
        else:
            return self.name