swingmusic-extended/app/lib/populate.py

from dataclasses import asdict
import os
from concurrent.futures import ProcessPoolExecutor
import platform

from requests import ConnectionError as RequestConnectionError
from requests import ReadTimeout

from app import settings
from app.lib.artistlib import CheckArtistImages
from app.lib.colorlib import ProcessAlbumColors, ProcessArtistColors
from app.lib.taglib import extract_thumb
from app.logger import log
from app.models import Album, Artist
from app.models.lastfm import SimilarArtist
from app.requests.artists import fetch_similar_artists
from app.store.albums import AlbumStore
from app.store.artists import ArtistStore
from app.utils.network import has_connection
from app.utils.progressbar import tqdm

from app.db.userdata import SimilarArtistTable


class CordinateMedia:
    """
    Cordinates the extracting of thumbnails
    """

    def __init__(self, instance_key: str):
        ProcessTrackThumbnails()
        ProcessAlbumColors()
        ProcessArtistColors()

        tried_to_download_new_images = False

        if has_connection():
            tried_to_download_new_images = True
            try:
                CheckArtistImages()
            except (RequestConnectionError, ReadTimeout) as e:
                log.error(
                    "Internet connection lost. Downloading artist images suspended."
                )
                log.error(e)  # REVIEW More informations = good
        else:
            log.warning("No internet connection. Downloading artist images suspended!")

        # Re-process the new artist images.
        if tried_to_download_new_images:
            ProcessArtistColors()

        if has_connection():
            print("Attempting to download similar artists...")
            FetchSimilarArtistsLastFM()


def get_image(album: Album):
    """
    The function retrieves an image from an album by iterating through its tracks and extracting the thumbnail from the first track that has one.

    :param album: An instance of the `Album` class representing the album to retrieve the image from.
    :type album: Album
    :return: None
    """
    matching_tracks = AlbumStore.get_album_tracks(album.albumhash)

    for track in matching_tracks:
        extracted = extract_thumb(track.filepath, track.albumhash + ".webp")

        if extracted:
            return


class ProcessTrackThumbnails:
    """
    Extracts the album art from all albums in album store.
    """

    def extract(self, albums: list[Album]):
        """
        Extracts the album art with platform specific logic.
        """

        if platform.system() == "Linux":
            # INFO: Processess are forked with access to global stores
            # It's "safe" to use a process pool
            cpus = max(1, os.cpu_count() // 2)
            with ProcessPoolExecutor(max_workers=cpus) as executor:
                results = list(
                    tqdm(
                        executor.map(get_image, albums),
                        total=len(albums),
                        desc="Extracting track images",
                    )
                )

                list(results)
        else:
            # INFO: Use a for loop for windows (and others I guess)
            for album in tqdm(albums, desc="Extracting track images"):
                get_image(album)

    def __init__(self) -> None:
        """
        Filters out albums that already have thumbnails and
        extracts the thumbnail for the other albums.
        """
        path = settings.Paths.get_sm_thumb_path()

        # read all the files in the thumbnail directory
        processed = set(i.replace(".webp", "") for i in os.listdir(path))
        # filter out albums that already have thumbnails
        albums = filter(
            lambda album: album.albumhash not in processed,
            AlbumStore.get_flat_list(),
        )

        albums = list(albums)
        self.extract(albums)


def save_similar_artists(artist: Artist):
    """
    Downloads and saves similar artists to the database.
    """
    if SimilarArtistTable.exists(artist.artisthash):
        return

    artists = fetch_similar_artists(artist.name)

    # INFO: Nones mean there was a connection error
    if artists is None:
        return

    artist_ = SimilarArtist(artist.artisthash, artists)
    SimilarArtistTable.insert_one(asdict(artist_))


class FetchSimilarArtistsLastFM:
    """
    Fetches similar artists from LastFM using a thread pool.
    """

    def __init__(self) -> None:
        # read all artists from db
        processed = set(a.artisthash for a in SimilarArtistTable.get_all())

        # filter out artists that already have similar artists
        artists = filter(
            lambda a: a.artisthash not in processed, ArtistStore.get_flat_list()
        )
        artists = list(artists)

        with ProcessPoolExecutor(max_workers=max(1, os.cpu_count() // 2)) as executor:
            try:
                print("Processing similar artists")
                results = list(
                    tqdm(
                        executor.map(save_similar_artists, artists),
                        total=len(artists),
                        desc="Fetching similar artists",
                    )
                )

                list(results)
            # any exception that can be raised by the pool
            except Exception as e:
                log.warn(e)
                return