music-kraken-core/music_kraken/objects/artwork.py

from __future__ import annotations

from copy import copy
from dataclasses import dataclass, field
from functools import cached_property
from typing import Dict, List, Optional, Set, Tuple, Type, TypedDict, Union

from ..connection import Connection
from ..utils import create_dataclass_instance, custom_hash
from ..utils.config import main_settings
from ..utils.enums import PictureType
from ..utils.string_processing import hash_url, unify
from .collection import Collection
from .metadata import ID3Timestamp
from .metadata import Mapping as id3Mapping
from .metadata import Metadata
from .parents import OuterProxy as Base
from .target import Target

artwork_connection: Connection = Connection(module="artwork")


@dataclass
class ArtworkVariant:
    url: str
    width: Optional[int] = None
    height: Optional[int] = None
    image_format: Optional[str] = None

    def __hash__(self) -> int:
        return custom_hash(self.url)

    def __eq__(self, other: ArtworkVariant) -> bool:
        return hash(self) == hash(other)

    def __contains__(self, other: str) -> bool:
        return custom_hash(other) == hash(self.url)

    def __merge__(self, other: ArtworkVariant) -> None:
        for key, value in other.__dict__.items():
            if value is None:
                continue

            if getattr(self, key) is None:
                setattr(self, key, value)

    @cached_property
    def target(self) -> Target:
        return Target.temp()

    def fetch(self) -> None:
        global artwork_connection

        r = artwork_connection.get(self.url, name=hash_url(url))
        if r is None:
            return

        self.target.raw_content = r.content

@dataclass
class Artwork:
    variants: List[ArtworkVariant] = field(default_factory=list)

    artwork_type: PictureType = PictureType.OTHER

    def search_variant(self, url: str) -> Optional[ArtworkVariant]:
        if url is None: 
            return None

        for variant in self.variants:
            if url in variant:
                return variant

        return None

    def __contains__(self, other: str) -> bool:
        return self.search_variant(other) is not None

    def add_data(self, **kwargs) -> None:
        variant = self.search_variant(kwargs.get("url"))

        if variant is None:
            variant, kwargs = create_dataclass_instance(ArtworkVariant, kwargs)
            self.variants.append(variant)

        variant.__dict__.update(kwargs)

    @property
    def url(self) -> Optional[str]:
        if len(self.variants) <= 0:
            return None
        return self.variants[0].url

    def fetch(self) -> None:
        for variant in self.variants:
            variant.fetch()


class ArtworkCollection:
    """
    Stores all the images/artworks for one data object.
    
    There could be duplicates before calling ArtworkCollection.compile()  
    _this is called before one object is downloaded automatically._
    """

    artwork_type: PictureType = PictureType.OTHER

    def __init__(
        self, 
        *data: List[Artwork], 
        parent_artworks: Set[ArtworkCollection] = None, 
        crop_images: bool = True,
    ) -> None:
        # this is used for the song artwork, to fall back to the song artwork
        self.parent_artworks: Set[ArtworkCollection] = parent_artworks or set()
        self.crop_images: bool = crop_images
        
        self._data = []
        self.extend(data)

    def search_artwork(self, url: str) -> Optional[ArtworkVariant]:
        for artwork in self._data:
            if url in artwork:
                return artwork

        return None
    
    def __contains__(self, other: str) -> bool:
        return self.search_artwork(other) is not None

    def _create_new_artwork(self, **kwargs) -> Tuple[Artwork, dict]:
        kwargs["artwork_type"] = kwargs.get("artwork_type", self.artwork_type)

        return create_dataclass_instance(Artwork, dict(**kwargs))

    def add_data(self, url: str, **kwargs) -> Artwork:
        kwargs["url"] = url

        artwork = self.search_artwork(url)

        if artwork is None:
            artwork, kwargs = self._create_new_artwork(**kwargs)
            self._data.append(artwork)

        artwork.add_data(**kwargs)
        return artwork

    def append(self, value: Union[Artwork, ArtworkVariant, dict], **kwargs):
        """
        You can append the types Artwork, ArtworkVariant or dict 
        the best option would be to use Artwork and avoid the other options.
        """
        if isinstance(value, dict):
            kwargs.update(value)
            value, kwargs = create_dataclass_instance(ArtworkVariant, kwargs)
        
        if isinstance(value, ArtworkVariant):
            kwargs["variants"] = [value]
            value, kwargs = create_dataclass_instance(Artwork, kwargs)

        if isinstance(value, Artwork):
            self._data.append(value)
            return
    
    def extend(self, values: List[Union[Artwork, ArtworkVariant, dict]], **kwargs):
        for value in values:
            self.append(value, **kwargs)

    def compile(self) -> None:
        """
        This will make the artworks ready for download
        """
        for artwork in self._data:
            artwork.fetch()

    def __merge__(self, other: ArtworkCollection, **kwargs) -> None:
        self.parent_artworks.update(other.parent_artworks)

        for key, value in other._variant_mapping.items():
            if key not in self._variant_mapping:
                self._variant_mapping[key] = value

    def __hash__(self) -> int:
        return id(self)

    def __iter__(self) -> Generator[Artwork, None, None]:
        yield from self._data

    def get_urls(self) -> Generator[str, None, None]:
        yield from (artwork.url for artwork in self._data if artwork.url is not None)

    """
    @property
    def flat_empty(self) -> bool:
        return len(self._variant_mapping.keys()) <= 0

    def _get_best_from_list(self, artwork_variants: List[ArtworkVariant]) -> Optional[ArtworkVariant]:
        return min(artwork_variants, key=lambda x: x["deviation"])

    @property
    def best_variant(self) -> ArtworkVariant:
        if self.flat_empty:
            return self._get_best_from_list([parent.best_variant for parent in self.parent_artworks])
        return self._get_best_from_list(self._variant_mapping.values())

    def get_variant_name(self, variant: ArtworkVariant) -> str:
        return f"artwork_{variant['width']}x{variant['height']}_{hash_url(variant['url']).replace('/', '_')}"
    """