StreamingCommunity 3.3.9__py3-none-any.whl → 3.4.2__py3-none-any.whl

StreamingCommunity/Lib/Downloader/DASH/segments.py

@@ -3,6 +3,7 @@
 import os
 import asyncio
 import time
+from typing import Dict, Optional
 
 
 # External libraries
@@ -24,6 +25,7 @@ DEFAULT_VIDEO_WORKERS = config_manager.get_int('M3U8_DOWNLOAD', 'default_video_w
 DEFAULT_AUDIO_WORKERS = config_manager.get_int('M3U8_DOWNLOAD', 'default_audio_workers')
 SEGMENT_MAX_TIMEOUT = config_manager.get_int("M3U8_DOWNLOAD", "segment_timeout")
 LIMIT_SEGMENT = config_manager.get_int('M3U8_DOWNLOAD', 'limit_segment')
+ENABLE_RETRY = config_manager.get_bool('M3U8_DOWNLOAD', 'enable_retry')
 
 
 # Variable
@@ -50,7 +52,8 @@ class MPD_Segments:
             self.limit_segments = LIMIT_SEGMENT if LIMIT_SEGMENT > 0 else None
         else:
             self.limit_segments = limit_segments
-            
+        
+        self.enable_retry = ENABLE_RETRY
         self.download_interrupted = False
         self.info_nFailed = 0
         
@@ -63,9 +66,12 @@ class MPD_Segments:
         self._last_progress_update = 0
         self._progress_update_interval = 0.1
         
-        # Segment tracking
-        self.segment_files = {}
+        # Segment tracking - store only metadata, not content
+        self.segment_status = {}  # {idx: {'downloaded': bool, 'size': int}}
         self.segments_lock = asyncio.Lock()
+        
+        # Estimator for progress tracking
+        self.estimator: Optional[M3U8_Ts_Estimator] = None
 
     def get_concat_path(self, output_dir: str = None):
         """
@@ -114,7 +120,7 @@ class MPD_Segments:
 
     async def download_segments(self, output_dir: str = None, concurrent_downloads: int = None, description: str = "DASH"):
         """
-        Download and concatenate all segments (including init) asynchronously and in order.
+        Download segments to temporary files, then concatenate them in order.
         
         Parameters:
             - output_dir (str): Output directory for segments
@@ -148,9 +154,9 @@ class MPD_Segments:
         semaphore = asyncio.Semaphore(concurrent_downloads)
 
         # Initialize estimator
-        estimator = M3U8_Ts_Estimator(total_segments=len(segment_urls) + 1)
+        self.estimator = M3U8_Ts_Estimator(total_segments=len(segment_urls) + 1)
 
-        self.segment_files = {}
+        self.segment_status = {}
         self.downloaded_segments = set()
         self.info_nFailed = 0
         self.download_interrupted = False
@@ -164,20 +170,21 @@ class MPD_Segments:
             async with httpx.AsyncClient(timeout=timeout_config, limits=limits) as client:
                 
                 # Download init segment
-                await self._download_init_segment(client, init_url, concat_path, estimator, progress_bar)
+                await self._download_init_segment(client, init_url, concat_path, progress_bar)
 
-                # Download all segments (first batch) - writes to temp files
+                # Download all segments to temp files
                 await self._download_segments_batch(
-                    client, segment_urls, temp_dir, semaphore, REQUEST_MAX_RETRY, estimator, progress_bar
+                    client, segment_urls, temp_dir, semaphore, REQUEST_MAX_RETRY, progress_bar
                 )
 
-                # Retry failed segments 
-                await self._retry_failed_segments(
-                    client, segment_urls, temp_dir, semaphore, REQUEST_MAX_RETRY, estimator, progress_bar
-                )
+                # Retry failed segments only if enabled
+                if self.enable_retry:
+                    await self._retry_failed_segments(
+                        client, segment_urls, temp_dir, semaphore, REQUEST_MAX_RETRY, progress_bar
+                    )
 
-                # Concatenate all segment files in order
-                await self._concatenate_segments(concat_path, len(segment_urls))
+                # Concatenate all segments IN ORDER
+                await self._concatenate_segments_in_order(temp_dir, concat_path, len(segment_urls))
 
         except KeyboardInterrupt:
             self.download_interrupted = True
@@ -189,7 +196,7 @@ class MPD_Segments:
         self._verify_download_completion()
         return self._generate_results(stream_type)
 
-    async def _download_init_segment(self, client, init_url, concat_path, estimator, progress_bar):
+    async def _download_init_segment(self, client, init_url, concat_path, progress_bar):
         """
         Download the init segment and update progress/estimator.
         """
@@ -205,49 +212,56 @@ class MPD_Segments:
             with open(concat_path, 'wb') as outfile:
                 if response.status_code == 200:
                     outfile.write(response.content)
-                    estimator.add_ts_file(len(response.content))
+                    if self.estimator:
+                        self.estimator.add_ts_file(len(response.content))
 
             progress_bar.update(1)
-            self._throttled_progress_update(len(response.content), estimator, progress_bar)
+            if self.estimator:
+                self._throttled_progress_update(len(response.content), progress_bar)
 
         except Exception as e:
             progress_bar.close()
             raise RuntimeError(f"Error downloading init segment: {e}")
 
-    def _throttled_progress_update(self, content_size: int, estimator, progress_bar):
+    def _throttled_progress_update(self, content_size: int, progress_bar):
         """
         Throttled progress update to reduce CPU usage.
         """
         current_time = time.time()
         if current_time - self._last_progress_update > self._progress_update_interval:
-            estimator.update_progress_bar(content_size, progress_bar)
+            if self.estimator:
+                self.estimator.update_progress_bar(content_size, progress_bar)
             self._last_progress_update = current_time
 
-    async def _download_segments_batch(self, client, segment_urls, temp_dir, semaphore, max_retry, estimator, progress_bar):
+    async def _download_segments_batch(self, client, segment_urls, temp_dir, semaphore, max_retry, progress_bar):
         """
-        Download a batch of segments and write them to temp files immediately.
+        Download segments to temporary files - write immediately to disk, not memory.
         """
         async def download_single(url, idx):
             async with semaphore:
                 headers = {'User-Agent': get_userAgent()}
+                temp_file = os.path.join(temp_dir, f"seg_{idx:06d}.tmp")
                 
                 for attempt in range(max_retry):
                     if self.download_interrupted:
-                        return idx, False, attempt
+                        return idx, False, attempt, 0
                         
                     try:
                         timeout = min(SEGMENT_MAX_TIMEOUT, 10 + attempt * 3)
                         resp = await client.get(url, headers=headers, follow_redirects=True, timeout=timeout)
 
-                        # Write to temp file immediately
+                        # Write directly to temp file
                         if resp.status_code == 200:
-                            temp_file = os.path.join(temp_dir, f"seg_{idx:06d}.tmp")
+                            content_size = len(resp.content)
+                            with open(temp_file, 'wb') as f:
+                                f.write(resp.content)
+                            
+                            # Update status
                             async with self.segments_lock:
-                                with open(temp_file, 'wb') as f:
-                                    f.write(resp.content)
-                                self.segment_files[idx] = temp_file
+                                self.segment_status[idx] = {'downloaded': True, 'size': content_size}
+                                self.downloaded_segments.add(idx)
                             
-                            return idx, True, attempt, len(resp.content)
+                            return idx, True, attempt, content_size
                         else:
                             if attempt < 2:
                                 sleep_time = 0.5 + attempt * 0.5
@@ -258,19 +272,21 @@ class MPD_Segments:
                     except Exception:
                         sleep_time = min(2.0, 1.1 * (2 ** attempt))
                         await asyncio.sleep(sleep_time)
+                
+                # Mark as failed
+                async with self.segments_lock:
+                    self.segment_status[idx] = {'downloaded': False, 'size': 0}
                         
                 return idx, False, max_retry, 0
 
-        # Initial download attempt
+        # Download all segments concurrently
         tasks = [download_single(url, i) for i, url in enumerate(segment_urls)]
 
         for coro in asyncio.as_completed(tasks):
             try:
                 idx, success, nretry, size = await coro
                 
-                if success:
-                    self.downloaded_segments.add(idx)
-                else:
+                if not success:
                     self.info_nFailed += 1
                 
                 if nretry > self.info_maxRetry:
@@ -278,15 +294,16 @@ class MPD_Segments:
                 self.info_nRetry += nretry
                     
                 progress_bar.update(1)
-                estimator.add_ts_file(size)
-                self._throttled_progress_update(size, estimator, progress_bar)
+                if self.estimator:
+                    self.estimator.add_ts_file(size)
+                    self._throttled_progress_update(size, progress_bar)
 
             except KeyboardInterrupt:
                 self.download_interrupted = True
-                print("\n[red]Download interrupted by user (Ctrl+C).")
+                console.print("\n[red]Download interrupted by user (Ctrl+C).")
                 break
 
-    async def _retry_failed_segments(self, client, segment_urls, temp_dir, semaphore, max_retry, estimator, progress_bar):
+    async def _retry_failed_segments(self, client, segment_urls, temp_dir, semaphore, max_retry, progress_bar):
         """
         Retry failed segments up to 3 times.
         """
@@ -301,6 +318,7 @@ class MPD_Segments:
             async def download_single(url, idx):
                 async with semaphore:
                     headers = {'User-Agent': get_userAgent()}
+                    temp_file = os.path.join(temp_dir, f"seg_{idx:06d}.tmp")
 
                     for attempt in range(max_retry):
                         if self.download_interrupted:
@@ -310,15 +328,17 @@ class MPD_Segments:
                             timeout = min(SEGMENT_MAX_TIMEOUT, 15 + attempt * 5)
                             resp = await client.get(url, headers=headers, timeout=timeout)
                             
-                            # Write to temp file immediately
+                            # Write directly to temp file
                             if resp.status_code == 200:
-                                temp_file = os.path.join(temp_dir, f"seg_{idx:06d}.tmp")
+                                content_size = len(resp.content)
+                                with open(temp_file, 'wb') as f:
+                                    f.write(resp.content)
+                                
                                 async with self.segments_lock:
-                                    with open(temp_file, 'wb') as f:
-                                        f.write(resp.content)
-                                    self.segment_files[idx] = temp_file
+                                    self.segment_status[idx] = {'downloaded': True, 'size': content_size}
+                                    self.downloaded_segments.add(idx)
                                 
-                                return idx, True, attempt, len(resp.content)
+                                return idx, True, attempt, content_size
                             else:
                                 await asyncio.sleep(1.5 * (2 ** attempt))
 
@@ -334,9 +354,7 @@ class MPD_Segments:
                 try:
                     idx, success, nretry, size = await coro
 
-                    if success:
-                        self.downloaded_segments.add(idx)
-                    else:
+                    if not success:
                         nFailed_this_round += 1
 
                     if nretry > self.info_maxRetry:
@@ -344,8 +362,9 @@ class MPD_Segments:
                     self.info_nRetry += nretry
                     
                     progress_bar.update(0)
-                    estimator.add_ts_file(size)
-                    self._throttled_progress_update(size, estimator, progress_bar)
+                    if self.estimator:
+                        self.estimator.add_ts_file(size)
+                        self._throttled_progress_update(size, progress_bar)
 
                 except KeyboardInterrupt:
                     self.download_interrupted = True
@@ -355,20 +374,24 @@ class MPD_Segments:
             self.info_nFailed = nFailed_this_round
             global_retry_count += 1
 
-    async def _concatenate_segments(self, concat_path, total_segments):
+    async def _concatenate_segments_in_order(self, temp_dir, concat_path, total_segments):
         """
-        Concatenate all segment files in order to the final output file.
-        Skip missing segments and continue with available ones.
+        Concatenate all segment files IN ORDER to the final output file.
         """
-        successful_segments = 0
         with open(concat_path, 'ab') as outfile:
             for idx in range(total_segments):
-                if idx in self.segment_files:
-                    temp_file = self.segment_files[idx]
-                    if os.path.exists(temp_file):
-                        with open(temp_file, 'rb') as infile:
-                            outfile.write(infile.read())
-                        successful_segments += 1
+                temp_file = os.path.join(temp_dir, f"seg_{idx:06d}.tmp")
+                
+                # Only concatenate successfully downloaded segments
+                if idx in self.downloaded_segments and os.path.exists(temp_file):
+                    with open(temp_file, 'rb') as infile:
+                        
+                        # Read and write in chunks to avoid memory issues
+                        while True:
+                            chunk = infile.read(8192)  # 8KB chunks
+                            if not chunk:
+                                break
+                            outfile.write(chunk)
 
     def _get_bar_format(self, description: str) -> str:
         """
@@ -435,7 +458,8 @@ class MPD_Segments:
         # Delete temp segment files
         if temp_dir and os.path.exists(temp_dir):
             try:
-                for temp_file in self.segment_files.values():
+                for idx in range(len(self.selected_representation.get('segment_urls', []))):
+                    temp_file = os.path.join(temp_dir, f"seg_{idx:06d}.tmp")
                     if os.path.exists(temp_file):
                         os.remove(temp_file)
                 os.rmdir(temp_dir)
@@ -447,7 +471,7 @@ class MPD_Segments:
             self._display_error_summary()
             
         # Clear memory
-        self.segment_files = {}
+        self.segment_status = {}
 
     def _display_error_summary(self) -> None:
         """
@@ -455,11 +479,28 @@ class MPD_Segments:
         """
         total_segments = len(self.selected_representation.get('segment_urls', []))
         failed_indices = [i for i in range(total_segments) if i not in self.downloaded_segments]
-        successful_segments = len(self.downloaded_segments)
-
-        console.print(f"[green]Download Summary: "
-              f"[cyan]Successful: [red]{successful_segments}/{total_segments} "
-              f"[cyan]Max retries: [red]{getattr(self, 'info_maxRetry', 0)} "
-              f"[cyan]Total retries: [red]{getattr(self, 'info_nRetry', 0)} "
-              f"[cyan]Failed segments: [red]{getattr(self, 'info_nFailed', 0)} "
-              f"[cyan]Failed indices: [red]{failed_indices} \n")
+
+        console.print(f" [cyan]Max retries: [red]{getattr(self, 'info_maxRetry', 0)} [white]| "
+            f"[cyan]Total retries: [red]{getattr(self, 'info_nRetry', 0)} [white]| "
+            f"[cyan]Failed segments: [red]{getattr(self, 'info_nFailed', 0)} [white]| "
+            f"[cyan]Failed indices: [red]{failed_indices}")
+    
+    def get_progress_data(self) -> Dict:
+        """Returns current download progress data for API."""
+        if not self.estimator:
+            return None
+            
+        total = self.get_segments_count()
+        downloaded = len(self.downloaded_segments)
+        percentage = (downloaded / total * 100) if total > 0 else 0
+        stats = self.estimator.get_stats(downloaded, total)
+        
+        return {
+            'total_segments': total,
+            'downloaded_segments': downloaded,
+            'failed_segments': self.info_nFailed,
+            'current_speed': stats['download_speed'],
+            'estimated_size': stats['estimated_total_size'],
+            'percentage': round(percentage, 2),
+            'eta_seconds': stats['eta_seconds']
+        }