Create TrackStreamInfo dataclass & Add proper docstrings

2 files changed, 52 insertions, 13 deletions

diff --git a/tidal_dl_ng/config.py b/tidal_dl_ng/config.py
index 134f3ff..8be4459 100644
--- a/tidal_dl_ng/config.py
+++ b/tidal_dl_ng/config.py
@@ -179,6 +179,9 @@ class Tidal(BaseConfig, metaclass=SingletonMeta):
         """
         Switches the shared session to Dolby Atmos credentials.
         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 self.is_atmos_session:
@@ -204,6 +207,12 @@ class Tidal(BaseConfig, metaclass=SingletonMeta):
         """
         Restores the shared session to the original user credentials.
         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 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_secret = self.original_client_secret
 
-        # Re-apply user's quality setting
-        self.settings_apply()
+        # Explicitly restore audio quality to user's configured setting
+        self.session.audio_quality = tidalapi.Quality(self.settings.data.quality_audio)
 
         # Re-login with original credentials
         if not self.login_token(do_pkce=self.is_pkce):
diff --git a/tidal_dl_ng/download.py b/tidal_dl_ng/download.py
index f5b7608..e07f463 100644
--- a/tidal_dl_ng/download.py
+++ b/tidal_dl_ng/download.py
@@ -16,6 +16,7 @@ import tempfile
 import time
 from collections.abc import Callable
 from concurrent import futures
+from dataclasses import dataclass
 from threading import Event
 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
 
 
+@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.
 # https://github.com/globocom/m3u8#using-different-http-clients
 class RequestsClient:
@@ -134,7 +145,8 @@ class Download:
         Args:
             tidal_obj (Tidal): TIDAL configuration object. Required for:
                 - 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.
             fn_logger (Callable): Logger function or object.
             skip_existing (bool, optional): Whether to skip existing files. Defaults to False.
@@ -789,7 +801,7 @@ class Download:
             media (Track | Video): Media item.
 
         Returns:
-            tuple[StreamManifest | None, str, bool, Stream | None]: Stream info.
+        tuple[StreamManifest | None, str, bool, Stream | None]: Stream info.
         """
         stream_manifest: StreamManifest | None = None
         media_stream: Stream | None = None
@@ -821,10 +833,16 @@ class Download:
         with self.tidal.stream_lock:
             try:
                 if isinstance(media, Track):
-                    stream_manifest, file_extension, do_flac_extract, media_stream = self._get_track_stream_info(media)
-                    if stream_manifest is None:
+                    track_info = self._get_track_stream_info(media)
+
+                    if track_info.stream_manifest is 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):
                     # Videos always require the 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
 
-    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.
-        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 = (
             self.settings.data.download_dolby_atmos
@@ -868,25 +893,30 @@ class Download:
         if want_atmos:
             if not self.tidal.switch_to_atmos_session():
                 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:
             if not self.tidal.restore_normal_session():
                 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()
 
         stream_manifest = media_stream.get_stream_manifest()
         file_extension = stream_manifest.file_extension
-        do_flac_extract = False
+        requires_flac_extraction = False
 
         if self.settings.data.extract_flac and (
             stream_manifest.codecs.upper() == Codec.FLAC and 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(
         self,