Wobuzz/wobuzz/player/track.py

#!/usr/bin/python3

from pydub import AudioSegment
from pygame.mixer import Sound
from tinytag import TinyTag
from tinytag.tinytag import Images as TTImages
from dataclasses import dataclass


@dataclass
class TrackMetadata:
    path: str | None=None
    title: str | None=None
    artist: str | None=None
    album: str | None=None
    genre: str | None=None
    images: TTImages | None=None  # tinytag images

    def add_missing(self):
        # Make the album be an empty string instead of "None"
        if self.title == "None":
            self.title = ""

        if self.artist == "None":
            self.artist = ""

        if self.album == "None":
            self.album = ""

        if self.genre == "None":
            self.genre = ""

        if self.path is None:  # can't add missing information without a path
            return

        if self.title is None or self.artist is None or self.album is None or self.genre is None:
            tags = TinyTag.get(self.path, ignore_errors=True, duration=False)

            self.title = tags.title
            self.artist = tags.artist
            self.album = tags.album
            self.genre = tags.genre


class Track:
    """
    Class representing a track.
    """

    def __init__(self, app, path: str, cache: bool=False, metadata: TrackMetadata=None):
        self.app = app
        self.path = path

        # add self to loaded tracks to make sure that no other track object is created for this track
        app.library.loaded_tracks[self.path] = self

        if metadata is None:
            # load metadata from audio file
            tags = TinyTag.get(path, ignore_errors=True, duration=False)

            self.metadata = TrackMetadata(path, tags.title, tags.artist, tags.album)

        else:
            self.metadata = metadata

        self.cached = False
        self.audio = None
        self.sound = None
        self.duration = 0

        self.items = []
        self.occurrences = {}  # all occurrences in playlists categorized by playlist and id of the track widget

        if cache:
            self.cache()

    def __new__(cls, app, path: str, cache: bool=False, metadata: TrackMetadata=None):
        loaded_track = app.library.loaded_track(path)

        if loaded_track is not None:
            if cache:
                loaded_track.cache()

            return loaded_track

        else:
            return super().__new__(cls)

    def set_occurrences(self):
        # set track item for every occurrence of track in a playlist

        new_occurrences = {}

        for item in self.items:
            # create dict of item: item.index (actually the id of the item bc. the item can't be used as key)
            playlist_occurrences = new_occurrences.get(item.playlist, {})
            playlist_occurrences[id(item)] = item.index

            new_occurrences[item.playlist] = playlist_occurrences

        self.correct_occurrences(new_occurrences)

        self.occurrences = new_occurrences

    def correct_occurrences(self, new_occurrences):
        """
        If this track is the currently playing track, and it gets moved, this corrects the current playlist index.
        """

        if self.app.player.current_playlist is not None:
            if self.app.player.current_playlist.current_track is self:
                for item in self.items:
                    if (
                        item.playlist in self.occurrences and
                        self.occurrences[item.playlist][id(item)] == self.app.player.current_playlist.current_track_index
                    ):
                        self.app.player.current_playlist.set_track(new_occurrences[item.playlist][id(item)])

    def cache(self):
        self.load_audio()
        # audio = normalize(audio)

        wav = self.audio.export(format="wav")

        self.sound = Sound(wav)

        self.duration = len(self.audio)  # track duration in milliseconds

        # metadata with images
        self.metadata.images = TinyTag.get(self.path, ignore_errors=True, duration=False, image=True).images

        self.cached = True

    def clear_cache(self):
        self.cached = False

        self.audio = None
        self.sound = None
        self.duration = 0

        self.metadata.images = None

    def load_audio(self):
        self.audio = AudioSegment.from_file(self.path)

    def remaining(self, position: int):
        remaining_audio = self.audio[position:]

        wav = remaining_audio.export(format="wav")

        sound = Sound(wav)

        # return the remaining part of the track's audio and the duration of the remaining part
        return sound, len(remaining_audio)

    def delete_items(self, playlist):
        """
        Deletes all QTreeWidgetItems that correspond to this track and the given playlist.
        """

        for item in self.items:
            if id(item) in self.occurrences[playlist]:
                self.items.remove(item)

        self.occurrences.pop(playlist)