Create TrackStreamInfo dataclass & Add proper docstrings

This commit is contained in:
Rikrdoga
2025-11-12 15:18:35 -06:00
parent af58b0a83d
commit 4f0026fada
2 changed files with 52 additions and 13 deletions
+11 -2
View File
@@ -179,6 +179,9 @@ class Tidal(BaseConfig, metaclass=SingletonMeta):
""" """
Switches the shared session to Dolby Atmos credentials. Switches the shared session to Dolby Atmos credentials.
Only re-authenticates if not already in Atmos mode. Only re-authenticates if not already in Atmos mode.
Returns:
bool: True if successful or already in Atmos mode, False otherwise.
""" """
# If we are already in Atmos mode, do nothing. # If we are already in Atmos mode, do nothing.
if self.is_atmos_session: if self.is_atmos_session:
@@ -204,6 +207,12 @@ class Tidal(BaseConfig, metaclass=SingletonMeta):
""" """
Restores the shared session to the original user credentials. Restores the shared session to the original user credentials.
Only re-authenticates if not already in Normal mode. Only re-authenticates if not already in Normal mode.
Args:
force: If True, forces restoration even if already in Normal mode.
Returns:
bool: True if successful or already in Normal mode, False otherwise.
""" """
# If we are already in Normal mode (and not forced), do nothing. # If we are already in Normal mode (and not forced), do nothing.
if not self.is_atmos_session and not force: if not self.is_atmos_session and not force:
@@ -213,8 +222,8 @@ class Tidal(BaseConfig, metaclass=SingletonMeta):
self.session.config.client_id = self.original_client_id self.session.config.client_id = self.original_client_id
self.session.config.client_secret = self.original_client_secret self.session.config.client_secret = self.original_client_secret
# Re-apply user's quality setting # Explicitly restore audio quality to user's configured setting
self.settings_apply() self.session.audio_quality = tidalapi.Quality(self.settings.data.quality_audio)
# Re-login with original credentials # Re-login with original credentials
if not self.login_token(do_pkce=self.is_pkce): if not self.login_token(do_pkce=self.is_pkce):
+41 -11
View File
@@ -16,6 +16,7 @@ import tempfile
import time import time
from collections.abc import Callable from collections.abc import Callable
from concurrent import futures from concurrent import futures
from dataclasses import dataclass
from threading import Event from threading import Event
from uuid import uuid4 from uuid import uuid4
@@ -75,6 +76,16 @@ from tidal_dl_ng.model.downloader import DownloadSegmentResult
from tidal_dl_ng.model.gui_data import ProgressBars from tidal_dl_ng.model.gui_data import ProgressBars
@dataclass
class TrackStreamInfo:
"""Container for track stream information."""
stream_manifest: StreamManifest | None
file_extension: str
requires_flac_extraction: bool
media_stream: Stream | None
# TODO: Set appropriate client string and use it for video download. # TODO: Set appropriate client string and use it for video download.
# https://github.com/globocom/m3u8#using-different-http-clients # https://github.com/globocom/m3u8#using-different-http-clients
class RequestsClient: class RequestsClient:
@@ -134,7 +145,8 @@ class Download:
Args: Args:
tidal_obj (Tidal): TIDAL configuration object. Required for: tidal_obj (Tidal): TIDAL configuration object. Required for:
- session: Main TIDAL API session - session: Main TIDAL API session
- atmos_session_context(): Dolby Atmos credential switching - switch_to_atmos_session(): Dolby Atmos credential switching
- restore_normal_session(): Restore original session credentials
path_base (str): Base path for downloads. path_base (str): Base path for downloads.
fn_logger (Callable): Logger function or object. fn_logger (Callable): Logger function or object.
skip_existing (bool, optional): Whether to skip existing files. Defaults to False. skip_existing (bool, optional): Whether to skip existing files. Defaults to False.
@@ -789,7 +801,7 @@ class Download:
media (Track | Video): Media item. media (Track | Video): Media item.
Returns: Returns:
tuple[StreamManifest | None, str, bool, Stream | None]: Stream info. tuple[StreamManifest | None, str, bool, Stream | None]: Stream info.
""" """
stream_manifest: StreamManifest | None = None stream_manifest: StreamManifest | None = None
media_stream: Stream | None = None media_stream: Stream | None = None
@@ -821,10 +833,16 @@ class Download:
with self.tidal.stream_lock: with self.tidal.stream_lock:
try: try:
if isinstance(media, Track): if isinstance(media, Track):
stream_manifest, file_extension, do_flac_extract, media_stream = self._get_track_stream_info(media) track_info = self._get_track_stream_info(media)
if stream_manifest is None:
if track_info.stream_manifest is None:
return None, "", False, None return None, "", False, None
stream_manifest = track_info.stream_manifest
file_extension = track_info.file_extension
do_flac_extract = track_info.requires_flac_extraction
media_stream = track_info.media_stream
elif isinstance(media, Video): elif isinstance(media, Video):
# Videos always require the normal session # Videos always require the normal session
if not self.tidal.restore_normal_session(): if not self.tidal.restore_normal_session():
@@ -854,10 +872,17 @@ class Download:
return stream_manifest, file_extension, do_flac_extract, media_stream return stream_manifest, file_extension, do_flac_extract, media_stream
def _get_track_stream_info(self, media: Track) -> tuple[StreamManifest | None, str, bool, Stream | None]: def _get_track_stream_info(self, media: Track) -> TrackStreamInfo:
""" """
Gets stream info for a Track, handling Atmos/Normal session switching. Gets stream info for a Track, handling Atmos/Normal session switching.
This is a helper for _get_stream_info to reduce complexity.
Args:
media: The track to get stream information for.
Returns:
TrackStreamInfo: Container with stream manifest, file extension,
FLAC extraction flag, and media stream object.
Returns TrackStreamInfo with None/empty values if fails.
""" """
want_atmos = ( want_atmos = (
self.settings.data.download_dolby_atmos self.settings.data.download_dolby_atmos
@@ -868,25 +893,30 @@ class Download:
if want_atmos: if want_atmos:
if not self.tidal.switch_to_atmos_session(): if not self.tidal.switch_to_atmos_session():
self.fn_logger.error(f"Failed to switch to Atmos session for track: {media.id}") self.fn_logger.error(f"Failed to switch to Atmos session for track: {media.id}")
return None, "", False, None return TrackStreamInfo(None, "", False, None)
else: else:
if not self.tidal.restore_normal_session(): if not self.tidal.restore_normal_session():
self.fn_logger.error(f"Failed to restore normal session for track: {media.id}") self.fn_logger.error(f"Failed to restore normal session for track: {media.id}")
return None, "", False, None return TrackStreamInfo(None, "", False, None)
media_stream = self.session.track(media.id).get_stream() if want_atmos else media.get_stream() media_stream = self.session.track(media.id).get_stream() if want_atmos else media.get_stream()
stream_manifest = media_stream.get_stream_manifest() stream_manifest = media_stream.get_stream_manifest()
file_extension = stream_manifest.file_extension file_extension = stream_manifest.file_extension
do_flac_extract = False requires_flac_extraction = False
if self.settings.data.extract_flac and ( if self.settings.data.extract_flac and (
stream_manifest.codecs.upper() == Codec.FLAC and file_extension != AudioExtensions.FLAC stream_manifest.codecs.upper() == Codec.FLAC and file_extension != AudioExtensions.FLAC
): ):
file_extension = AudioExtensions.FLAC file_extension = AudioExtensions.FLAC
do_flac_extract = True requires_flac_extraction = True
return stream_manifest, file_extension, do_flac_extract, media_stream return TrackStreamInfo(
stream_manifest=stream_manifest,
file_extension=file_extension,
requires_flac_extraction=requires_flac_extraction,
media_stream=media_stream,
)
def _perform_actual_download( def _perform_actual_download(
self, self,