Update app.py

app.py

@@ -538,7 +538,13 @@ class DownloadManager:
             logger.error(f"Error getting sublinks: {e}")
             return []
 
+def safe_rerun():
+    """Safely rerun the app if experimental_rerun is available."""
+    if hasattr(st, "experimental_rerun"):
+        st.experimental_rerun()
+
 def main():
+    # Initialize session state on first run
     if 'initialized' not in st.session_state:
         st.session_state.initialized = True
         st.session_state.discovered_files = []
@@ -548,18 +554,11 @@ def main():
 
     st.title("Advanced File Downloader")
     
-    # Sidebar
+    # Sidebar configuration
     with st.sidebar:
-        # Mode Selection
         mode = st.radio("Select Mode", ["Manual URL", "Bing Search", "PDF Summarizer"], key="mode_select")
-        
-        # Advanced Options
         with st.expander("Advanced Options", expanded=True):
-            custom_extensions = st.text_input(
-                "Custom File Extensions",
-                placeholder=".csv, .txt, .epub",
-                key="custom_ext_input"
-            )
+            custom_extensions = st.text_input("Custom File Extensions", placeholder=".csv, .txt, .epub", key="custom_ext_input")
             max_sublinks = st.number_input(
                 "Maximum Sublinks to Process",
                 min_value=1,
@@ -580,19 +579,16 @@ def main():
             )
             use_proxy = st.checkbox("Use Proxy", key="proxy_checkbox")
             proxy = st.text_input("Proxy URL", placeholder="http://proxy:port", key="proxy_input")
-
-        # Google Drive Integration
         with st.expander("Google Drive Integration", expanded=False):
             if st.button("Start Google Sign-In", key="google_signin_btn"):
                 auth_url = get_google_auth_url()
                 st.markdown(f"[Click here to authorize]({auth_url})")
-                
             auth_code = st.text_input("Enter authorization code", key="auth_code_input")
             if st.button("Complete Sign-In", key="complete_signin_btn") and auth_code:
                 creds, msg = exchange_code_for_credentials(auth_code)
                 st.session_state.google_creds = creds
                 st.write(msg)
-
+    
     # Main content area
     if mode == "Manual URL":
         st.header("Manual URL Mode")
@@ -623,17 +619,16 @@ def main():
                     if files:
                         st.success(f"Found {len(files)} files!")
                         
-                        # Select All/Clear Selection buttons
+                        # File selection block 1
                         col1, col2 = st.columns([1, 4])
                         with col1:
                             if st.button("Select All", key="select_all_btn"):
                                 st.session_state.selected_files = list(range(len(files)))
-                                st.experimental_rerun()
+                                safe_rerun()
                             if st.button("Clear Selection", key="clear_selection_btn"):
                                 st.session_state.selected_files = []
-                                st.experimental_rerun()
+                                safe_rerun()
                         
-                        # File selection
                         selected_files = st.multiselect(
                             "Select files to download",
                             options=list(range(len(files))),
@@ -641,8 +636,6 @@ def main():
                             format_func=lambda x: f"{files[x]['filename']} ({files[x]['size']})",
                             key="file_multiselect"
                         )
-                        
-                        # Update session state
                         st.session_state.selected_files = selected_files
                         
                         if selected_files:
@@ -659,28 +652,20 @@ def main():
                             if st.button("Download Selected", key="download_btn"):
                                 if not os.path.exists(download_dir):
                                     os.makedirs(download_dir)
-                                    
+                                
                                 async def download_files():
                                     downloaded_paths = []
                                     progress_bar = st.progress(0)
                                     status_text = st.empty()
-                                    
                                     async with DownloadManager(use_proxy=use_proxy, proxy=proxy) as dm:
                                         for i, idx in enumerate(selected_files):
                                             progress = (i + 1) / len(selected_files)
                                             file_info = files[idx]
-                                            
                                             status_text.text(f"Downloading {file_info['filename']}... ({i+1}/{len(selected_files)})")
                                             progress_bar.progress(progress)
-                                            
-                                            path = await dm.download_file(
-                                                file_info,
-                                                download_dir,
-                                                url
-                                            )
+                                            path = await dm.download_file(file_info, download_dir, url)
                                             if path:
                                                 downloaded_paths.append(path)
-                                        
                                         status_text.empty()
                                         progress_bar.empty()
                                         return downloaded_paths
@@ -689,11 +674,9 @@ def main():
                                 
                                 if downloaded:
                                     st.success(f"Successfully downloaded {len(downloaded)} files")
-                                    
                                     if create_zip or upload_to_drive:
                                         zip_path = create_zip_file(downloaded, download_dir)
                                         st.success(f"Created ZIP file: {zip_path}")
-                                        
                                         if upload_to_drive and st.session_state.get('google_creds'):
                                             with st.spinner("Uploading to Google Drive..."):
                                                 drive_id = google_drive_upload(zip_path, st.session_state.google_creds)
@@ -701,7 +684,6 @@ def main():
                                                     st.success(f"Uploaded to Google Drive. File ID: {drive_id}")
                                                 else:
                                                     st.error(drive_id)
-                                        
                                         if delete_after:
                                             for path in downloaded:
                                                 try:
@@ -711,23 +693,19 @@ def main():
                                             st.info("Deleted original files after ZIP creation")
                     else:
                         st.warning("No files found.")
-
-        # Display current files if they exist in session state
+        
+        # File selection block 2 (if files are already discovered)
         if st.session_state.discovered_files:
             files = st.session_state.discovered_files
             st.success(f"Found {len(files)} files!")
-            
-            # Select All/Clear Selection buttons
             col1, col2 = st.columns([1, 4])
             with col1:
                 if st.button("Select All", key="select_all_btn2"):
                     st.session_state.selected_files = list(range(len(files)))
-                    st.experimental_rerun()
+                    safe_rerun()
                 if st.button("Clear Selection", key="clear_selection_btn2"):
                     st.session_state.selected_files = []
-                    st.experimental_rerun()
-            
-            # File selection
+                    safe_rerun()
             selected_files = st.multiselect(
                 "Select files to download",
                 options=list(range(len(files))),
@@ -735,10 +713,7 @@ def main():
                 format_func=lambda x: f"{files[x]['filename']} ({files[x]['size']})",
                 key="file_multiselect2"
             )
-            
-            # Update session state
             st.session_state.selected_files = selected_files
-            
             if selected_files:
                 col1, col2, col3, col4 = st.columns(4)
                 with col1:
@@ -749,45 +724,32 @@ def main():
                     delete_after = st.checkbox("Delete after creating ZIP", key="delete_after_checkbox2")
                 with col4:
                     upload_to_drive = st.checkbox("Upload to Google Drive", key="upload_drive_checkbox2")
-                
                 if st.button("Download Selected", key="download_btn2"):
                     if not os.path.exists(download_dir):
                         os.makedirs(download_dir)
-                        
+                    
                     async def download_files():
                         downloaded_paths = []
                         progress_bar = st.progress(0)
                         status_text = st.empty()
-                        
                         async with DownloadManager(use_proxy=use_proxy, proxy=proxy) as dm:
                             for i, idx in enumerate(selected_files):
                                 progress = (i + 1) / len(selected_files)
                                 file_info = files[idx]
-                                
                                 status_text.text(f"Downloading {file_info['filename']}... ({i+1}/{len(selected_files)})")
                                 progress_bar.progress(progress)
-                                
-                                path = await dm.download_file(
-                                    file_info,
-                                    download_dir,
-                                    st.session_state.current_url
-                                )
+                                path = await dm.download_file(file_info, download_dir, st.session_state.current_url)
                                 if path:
                                     downloaded_paths.append(path)
-                            
                             status_text.empty()
                             progress_bar.empty()
                             return downloaded_paths
-                    
                     downloaded = asyncio.run(download_files())
-                    
                     if downloaded:
                         st.success(f"Successfully downloaded {len(downloaded)} files")
-                        
                         if create_zip or upload_to_drive:
                             zip_path = create_zip_file(downloaded, download_dir)
                             st.success(f"Created ZIP file: {zip_path}")
-                            
                             if upload_to_drive and st.session_state.get('google_creds'):
                                 with st.spinner("Uploading to Google Drive..."):
                                     drive_id = google_drive_upload(zip_path, st.session_state.google_creds)
@@ -795,7 +757,6 @@ def main():
                                         st.success(f"Uploaded to Google Drive. File ID: {drive_id}")
                                     else:
                                         st.error(drive_id)
-                            
                             if delete_after:
                                 for path in downloaded:
                                     try:
@@ -803,12 +764,11 @@ def main():
                                     except Exception as e:
                                         st.warning(f"Could not delete {path}: {e}")
                                 st.info("Deleted original files after ZIP creation")
-
+    
     elif mode == "Bing Search":
         st.header("Bing Search Mode")
         query = st.text_input("Enter search query", key="search_query_input")
         num_results = st.slider("Number of results", 1, 50, 5, key="num_results_slider")
-        
         if st.button("Search", key="search_btn"):
             if query:
                 async def run_search():
@@ -835,21 +795,19 @@ def main():
                                                 st.session_state.discovered_files = files
                                                 st.session_state.current_url = url
                                                 st.session_state.selected_files = []
-                                                st.experimental_rerun()
+                                                safe_rerun()
                                             else:
                                                 st.warning("No files found on this page.")
                             else:
                                 st.warning("No search results found.")
-                
                 asyncio.run(run_search())
-
+    
     else:  # PDF Summarizer mode
         if summarizer is None:
             st.error("PDF summarization is not available due to model loading errors.")
         else:
             st.header("PDF Summarizer")
             pdf_url = st.text_input("Enter PDF URL", key="pdf_url_input")
-            
             if st.button("Summarize", key="summarize_btn"):
                 if pdf_url:
                     with st.spinner("Generating summary..."):
@@ -873,4 +831,3 @@ if __name__ == "__main__":
         main()
     except Exception as e:
         st.error(f"An error occurred: {str(e)}")
-        logger.error(f"Application error: {str(e)}", exc_info=True)