BrickTracker/bricktracker/rebrickable_image.py

import os
from typing import TYPE_CHECKING
from urllib.parse import urlparse

from flask import current_app, url_for
import requests
from shutil import copyfileobj

from .exceptions import DownloadException
if TYPE_CHECKING:
    from .minifigure import BrickMinifigure
    from .part import BrickPart
    from .set import BrickSet


# A set, part or minifigure image from Rebrickable
class RebrickableImage(object):
    brickset: 'BrickSet'
    minifigure: 'BrickMinifigure | None'
    part: 'BrickPart | None'

    extension: str | None

    def __init__(
        self,
        brickset: 'BrickSet',
        /,
        minifigure: 'BrickMinifigure | None' = None,
        part: 'BrickPart | None' = None,
    ):
        # Save all objects
        self.brickset = brickset
        self.minifigure = minifigure
        self.part = part

        # Currently everything is saved as 'jpg'
        self.extension = 'jpg'

        # Guess the extension
        # url = self.url()
        # if url is not None:
        #     _, extension = os.path.splitext(url)
        #     # TODO: Add allowed extensions
        #     if extension != '':
        #         self.extension = extension

    # Import the image from Rebrickable
    def download(self, /) -> None:
        path = self.path()

        # Avoid doing anything if the file exists
        if os.path.exists(path):
            return

        url = self.url()
        if url is None:
            return

        # Grab the image
        response = requests.get(url, stream=True)
        if response.ok:
            with open(path, 'wb') as f:
                copyfileobj(response.raw, f)
        else:
            raise DownloadException('could not get image {id} at {url}'.format(
                id=self.id(),
                url=url,
            ))

    # Return the folder depending on the objects provided
    def folder(self, /) -> str:
        if self.part is not None:
            return current_app.config['PARTS_FOLDER'].value

        if self.minifigure is not None:
            return current_app.config['MINIFIGURES_FOLDER'].value

        return current_app.config['SETS_FOLDER'].value

    # Return the id depending on the objects provided
    def id(self, /) -> str:
        if self.part is not None:
            if self.part.fields.part_img_url_id is None:
                return RebrickableImage.nil_name()
            else:
                return self.part.fields.part_img_url_id

        if self.minifigure is not None:
            if self.minifigure.fields.set_img_url is None:
                return RebrickableImage.nil_minifigure_name()
            else:
                return self.minifigure.fields.fig_num

        return self.brickset.fields.set_num

    # Return the path depending on the objects provided
    def path(self, /) -> str:
        return os.path.join(
            current_app.static_folder,  # type: ignore
            self.folder(),
            '{id}.{ext}'.format(id=self.id(), ext=self.extension),
        )

    # Return the url depending on the objects provided
    def url(self, /) -> str:
        if self.part is not None:
            if self.part.fields.part_img_url is None:
                return current_app.config['REBRICKABLE_IMAGE_NIL'].value
            else:
                return self.part.fields.part_img_url

        if self.minifigure is not None:
            if self.minifigure.fields.set_img_url is None:
                return current_app.config['REBRICKABLE_IMAGE_NIL_MINIFIGURE'].value  # noqa: E501
            else:
                return self.minifigure.fields.set_img_url

        return self.brickset.fields.set_img_url

    # Return the name of the nil image file
    @staticmethod
    def nil_name() -> str:
        filename, _ = os.path.splitext(
            os.path.basename(
                urlparse(current_app.config['REBRICKABLE_IMAGE_NIL'].value).path  # noqa: E501
            )
        )

        return filename

    # Return the name of the nil minifigure image file
    @staticmethod
    def nil_minifigure_name() -> str:
        filename, _ = os.path.splitext(
            os.path.basename(
                urlparse(current_app.config['REBRICKABLE_IMAGE_NIL_MINIFIGURE'].value).path  # noqa: E501
            )
        )

        return filename

    # Return the static URL for an image given a name and folder
    @staticmethod
    def static_url(name: str, folder_name: str) -> str:
        folder: str = current_app.config[folder_name].value

        # /!\ Everything is saved as .jpg, even if it came from a .png
        # not changing this behaviour.

        # Grab the extension
        # _, extension = os.path.splitext(self.part_img_url)
        extension = '.jpg'

        # Compute the path
        path = os.path.join(folder, '{name}{ext}'.format(
            name=name,
            ext=extension,
        ))

        return url_for('static', filename=path)