Update app.py

app.py

@@ -3,7 +3,7 @@ from pytubefix import YouTube
 import tempfile
 import os
 from urllib.error import HTTPError
-from typing import Optional, List, Dict
+from typing import Optional, List, Dict, Tuple
 import logging
 
 # Configure logging
@@ -33,29 +33,35 @@ class VideoDownloader:
             st.error(f"Error accessing video: {str(e)}")
             return None
 
-    def get_streams(self, yt: YouTube, stream_type: str = 'video') -> Dict:
-        """Get available streams based on type"""
+    def process_streams(self, yt: YouTube, stream_type: str = 'video') -> Tuple[List, List, str]:
+        """Process and return streams information"""
         try:
             if stream_type == 'video':
-                streams = yt.streams.filter(progressive=True, file_extension='mp4').all()
-                resolutions = sorted(list(set(stream.resolution for stream in streams if stream.resolution)), reverse=True)
-                return {
-                    'streams': streams,
-                    'qualities': resolutions,
-                    'title': yt.title
-                }
-            else:  # audio
-                streams = yt.streams.filter(only_audio=True).all()
-                bitrates = sorted(list(set(stream.abr for stream in streams if stream.abr)), reverse=True)
-                return {
-                    'streams': streams,
-                    'qualities': bitrates,
-                    'title': yt.title
-                }
+                # Get all progressive MP4 streams
+                all_streams = yt.streams.filter(progressive=True, file_extension='mp4').all()
+                # Extract unique resolutions
+                resolutions = []
+                seen = set()
+                for stream in all_streams:
+                    if stream.resolution and stream.resolution not in seen:
+                        resolutions.append(stream.resolution)
+                        seen.add(stream.resolution)
+                return all_streams, sorted(resolutions, reverse=True), yt.title
+            else:
+                # Get all audio streams
+                all_streams = yt.streams.filter(only_audio=True).all()
+                # Extract unique bitrates
+                bitrates = []
+                seen = set()
+                for stream in all_streams:
+                    if stream.abr and stream.abr not in seen:
+                        bitrates.append(stream.abr)
+                        seen.add(stream.abr)
+                return all_streams, sorted(bitrates, reverse=True), yt.title
         except Exception as e:
-            logger.error(f"Error getting streams: {e}")
-            st.error(f"Error getting available qualities: {str(e)}")
-            return {'streams': [], 'qualities': [], 'title': ''}
+            logger.error(f"Error processing streams: {e}")
+            st.error(f"Error processing available qualities: {str(e)}")
+            return [], [], ""
 
     def download_content(self, stream, title: str, extension: str) -> None:
         """Download and provide content to user"""
@@ -109,23 +115,23 @@ def main():
         if url:
             yt = downloader.get_youtube(url)
             if yt:
-                stream_data = downloader.get_streams(yt, 'video')
-                if stream_data:
-                    st.success(f"📺 Found: {stream_data['title']}")
+                streams, qualities, title = downloader.process_streams(yt, 'video')
+                if streams and qualities:
+                    st.success(f"📺 Found: {title}")
                     quality = st.selectbox(
                         "Select Video Quality:",
-                        stream_data['qualities']
+                        qualities
                     )
                     
                     selected_stream = next(
-                        (s for s in stream_data['streams'] if s.resolution == quality),
+                        (s for s in streams if s.resolution == quality),
                         None
                     )
                     
                     if selected_stream and st.button("Prepare Download", type="primary"):
                         downloader.download_content(
                             selected_stream,
-                            stream_data['title'],
+                            title,
                             'mp4'
                         )
     
@@ -137,23 +143,23 @@ def main():
         if url:
             yt = downloader.get_youtube(url)
             if yt:
-                stream_data = downloader.get_streams(yt, 'audio')
-                if stream_data:
-                    st.success(f"🎵 Found: {stream_data['title']}")
+                streams, qualities, title = downloader.process_streams(yt, 'audio')
+                if streams and qualities:
+                    st.success(f"🎵 Found: {title}")
                     quality = st.selectbox(
                         "Select Audio Quality:",
-                        stream_data['qualities']
+                        qualities
                     )
                     
                     selected_stream = next(
-                        (s for s in stream_data['streams'] if s.abr == quality),
+                        (s for s in streams if s.abr == quality),
                         None
                     )
                     
                     if selected_stream and st.button("Prepare Audio Download", type="primary"):
                         downloader.download_content(
                             selected_stream,
-                            stream_data['title'],
+                            title,
                             'mp3'
                         )