cbwebreader/comic/processing.py

from itertools import chain
from pathlib import Path

from django.conf import settings
from django.contrib.auth.base_user import AbstractBaseUser
from django.db.models import (
    Case,
    Count,
    ExpressionWrapper,
    F,
    IntegerField,
    PositiveSmallIntegerField,
    Q,
    QuerySet,
    When,
)

from comic import models
from comic.archives import get_archive_files
from comic.errors import NotCompatibleArchive


def generate_directory(user: AbstractBaseUser, directory: models.Directory | None = None) \
        -> list[models.Directory | models.ComicBook]:
    dir_path = Path(settings.COMIC_BOOK_VOLUME, directory.path) if directory else settings.COMIC_BOOK_VOLUME
    files = []

    dir_db_query = models.Directory.objects.filter(parent=directory)
    clean_directories(dir_db_query, dir_path, directory)

    file_db_query = models.ComicBook.objects.filter(directory=directory)
    clean_files(file_db_query, user, dir_path, directory)

    dir_db_query = dir_db_query.annotate(
        total=Count('comicbook', distinct=True),
        progress=Count('comicbook__comicstatus', Q(comicbook__comicstatus__finished=True,
                                                   comicbook__comicstatus__user=user), distinct=True),
        finished=ExpressionWrapper(Q(total=F('progress')), output_field=IntegerField()),
        unread=ExpressionWrapper(Q(total__gt=F('progress')), output_field=IntegerField()),
    )
    files.extend(dir_db_query)

    # Create Missing Status
    new_status = [models.ComicStatus(comic=file, user=user) for file in
                  file_db_query.exclude(comicstatus__in=models.ComicStatus.objects.filter(
                      comic__in=file_db_query, user=user))]
    models.ComicStatus.objects.bulk_create(new_status)

    file_db_query = file_db_query.annotate(
        progress=F('comicstatus__last_read_page') + 1,
        finished=F('comicstatus__finished'),
        unread=F('comicstatus__unread'),
        user=F('comicstatus__user'),
        classification=Case(
            When(directory__isnull=True, then=models.Directory.Classification.C_G),
            default=F('directory__classification'),
            output_field=PositiveSmallIntegerField(choices=models.Directory.Classification.choices)
        )
    ).filter(Q(user__isnull=True) | Q(user=user.id))

    files.extend(file_db_query)

    for file in chain(file_db_query, dir_db_query):
        if file.thumbnail and not Path(file.thumbnail.path).exists():
            file.thumbnail.delete()
            file.save()
    files.sort(key=lambda x: x.title)
    files.sort(key=lambda x: x.type, reverse=True)
    return files


def clean_directories(directories: QuerySet, dir_path: Path, directory: models.Directory | None = None) -> None:
    dir_db_set = {Path(settings.COMIC_BOOK_VOLUME, x.path) for x in directories}
    dir_list = {x for x in sorted(dir_path.glob('*')) if x.is_dir()}
    # Create new directories db instances
    for new_directory in dir_list - dir_db_set:
        models.Directory(name=new_directory.name, parent=directory).save()

    # Remove stale db instances
    for stale_directory in dir_db_set - dir_list:
        models.Directory.objects.get(name=stale_directory.name, parent=directory).delete()


def clean_files(files: QuerySet, user: AbstractBaseUser, dir_path: Path, directory: models.Directory | None = None) \
        -> None:
    file_list = {x for x in sorted(dir_path.glob('*')) if x.is_file()}
    files_db_set = {Path(dir_path, x.file_name) for x in files}

    # Parse new comics
    books_to_add = []
    for new_comic in file_list - files_db_set:
        if new_comic.suffix.lower() in settings.SUPPORTED_FILES:
            new_book = models.ComicBook(file_name=new_comic.name, directory=directory)
            try:
                archive, archive_type = new_book.get_archive()
                if archive_type == 'archive':
                    new_book.page_count = len(get_archive_files(archive))
                elif archive_type == 'pdf':
                    new_book.page_count = archive.page_count
            except NotCompatibleArchive:
                # Unreadable/corrupt archive - index the file with no page count
                pass
            books_to_add.append(new_book)
    models.ComicBook.objects.bulk_create(books_to_add)

    status_to_add = []
    for book in books_to_add:
        status_to_add.append(models.ComicStatus(user=user, comic=book))

    models.ComicStatus.objects.bulk_create(status_to_add)

    # Remove stale comic instances
    for stale_comic in files_db_set - file_list:
        models.ComicBook.objects.get(file_name=stale_comic.name, directory=directory).delete()