Enhance Download Functionality with Progress Bar and Improved Logging

requirements.txt

@@ -1,2 +1,3 @@
 clipboard==0.0.4
 requests==2.28.1
+tqdm==4.67.1

src/TSRDownload.py

@@ -2,6 +2,7 @@
 import requests, time, os, re
 from TSRUrl import TSRUrl
 from logger import logger
+from tqdm import tqdm
 from exceptions import *
 
 
@@ -17,44 +18,50 @@ def __init__(self, url: TSRUrl, sessionId: str):
 
     @classmethod
     def download(self, downloadPath: str) -> str:
-        logger.info(f"Starting download for: {self.url.url}")
+        logger.info(f"\n\n=== Starting download for: {self.url.url} ===")
         timeToSleep = 15000 - (time.time() * 1000 - self.ticketInitializedTime)
         if timeToSleep > 0:
+            logger.info(f"Waiting for {timeToSleep / 1000:.2f} seconds before starting...")
             time.sleep(timeToSleep / 1000)
 
         downloadUrl = self.__getDownloadUrl()
-        logger.debug(f"Got downloadUrl: {downloadUrl}")
+        logger.debug(f"Got download URL: {downloadUrl}")
         fileName = self.__getFileName(downloadUrl)
-        logger.debug(f"Got fileName: {fileName}")
+        logger.debug(f"Got file name: {fileName}")
 
         startingBytes = (
             os.path.getsize(f"{downloadPath}/{fileName}.part")
             if os.path.exists(f"{downloadPath}/{fileName}.part")
             else 0
         )
-        logger.debug(f"Got startingBytes: {startingBytes}")
+        logger.debug(f"Starting bytes: {startingBytes}")
         request = self.session.get(
             downloadUrl,
             stream=True,
             headers={"Range": f"bytes={startingBytes}-"},
         )
-        logger.debug(f"Request status is: {request.status_code}")
-        file = open(f"{downloadPath}/{fileName}.part", "wb")
+        logger.debug(f"Request status: {request.status_code}")
 
-        for index, chunk in enumerate(request.iter_content(1024 * 128)):
-            logger.debug(f"Downloading chunk #{index} of {downloadUrl}")
-            file.write(chunk)
-        file.close()
-        logger.debug(f"Removing .part from file name: {fileName}")
+        total_size = int(request.headers.get('content-length', 0)) + startingBytes  # Total size of the file
+        with open(f"{downloadPath}/{fileName}.part", "wb") as file:
+            # Initialize tqdm progress bar
+            with tqdm(total=total_size, initial=startingBytes, unit='B', unit_scale=True, desc=fileName) as bar:
+                for index, chunk in enumerate(request.iter_content(1024 * 128)):
+                    file.write(chunk)
+                    bar.update(len(chunk))  # Update progress bar
+
+        logger.debug(f"\nRenaming .part file to: {fileName}")
         if os.path.exists(f"{downloadPath}/{fileName}"):
-            logger.debug(f"{downloadPath}/{fileName} Already exists! Replacing file")
+            logger.debug(f"{downloadPath}/{fileName} already exists! Replacing file.")
             os.replace(f"{downloadPath}/{fileName}.part", f"{downloadPath}/{fileName}")
         else:
-            logger.debug(f"{downloadPath}/{fileName} doesn't exist! Renaming file")
+            logger.debug(f"{downloadPath}/{fileName} doesn't exist! Renaming file.")
             os.rename(
                 f"{downloadPath}/{fileName}.part",
                 f"{downloadPath}/{fileName}",
             )
+
+        logger.info(f"\n=== Download completed: {fileName} ===\n")
         return fileName
 
     @classmethod
@@ -88,4 +95,4 @@ def __getTSRDLTicketCookie(self) -> str:
         )
         self.session.get(self.url.downloadUrl)
         self.ticketInitializedTime = time.time() * 1000
-        return response.cookies.get("tsrdlticket")
+        return response.cookies.get("tsrdlticket")