Delete audio file after download

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 from pytube import YouTube
 import os
 import base64
+import time
 
 def get_binary_file_downloader_html(bin_file, file_label='File'):
     with open(bin_file, 'rb') as f:
@@ -18,6 +19,14 @@ def on_complete(stream, filename):
     print('stream:', stream)
     print('filename:', filename)
 
+def delete_file(filepath):
+    if os.path.exists(filepath):
+        os.remove(filepath)
+
+def on_download_file_click(filepath):
+    # sleep Xsec can delete file
+    time.sleep(2)
+    delete_file(filepath)
 
 def main():
     # Set the title and background color
@@ -28,7 +37,7 @@ def main():
     # st.markdown('<style>h3{color: pink;  text-align: center;}</style>', unsafe_allow_html=True)
 
     # Expander for app details
-    with st.expander("About the App"):
+    with st.expander("About the App (Works well on chrome browser)"):
         st.write("This app allows you to download MP3 audio from  YouTube video url.")
         st.write("Enter a YouTube URL in the input box below and click on 'Submit'")
 
@@ -44,13 +53,14 @@ def main():
         # extract only audio
         video = yt.streams.filter(only_audio=True).first()
 
+        initial_filename = yt.title
         # download the file
         destination = '.' # current folder
-        out_file = video.download(output_path=destination)
+        out_file = video.download(output_path=destination, filename=str(initial_filename))
 
         # save the file
         base, ext = os.path.splitext(out_file)
-        new_file = 'audio_file' + '.mp3' # base[0:10] + '.mp3'
+        new_file = '-'.join(initial_filename.split(' '))[-15:] + '.mp3' # base[0:10] + '.mp3'
         os.rename(out_file, new_file)
 
         # Display layout with 2 rows
@@ -63,8 +73,8 @@ def main():
 
         # download audio file
         with open(new_file, 'rb') as f:
-            st.download_button('Download Audio', f, file_name=new_file)
-            
+            st.download_button('Download Audio', f, file_name=new_file, on_click=on_download_file_click(new_file))
+                
         progress_bar.progress(100, text='DONE')
         st.success('Successfull!!')