From 3a99e0195bccd54f0b633b7751983689667e35ad Mon Sep 17 00:00:00 2001
From: Yannik Schmidt <yannik.schmidt@potaris.de>
Date: Sat, 7 Jun 2025 13:42:18 +0200
Subject: [PATCH] feat: implement progress for extraction operation

---
 db.py          |  1 +
 software.py    | 18 ++++++++++--------
 statekeeper.py | 38 ++++++++++++++++++++++++++++++++------
 3 files changed, 43 insertions(+), 14 deletions(-)

diff --git a/db.py b/db.py
index 6a4b5c0..69bf313 100644
--- a/db.py
+++ b/db.py
@@ -15,6 +15,7 @@ class Download(Base):
     local_path = Column(String)
     url = Column(String)
     size = Column(Integer)
+    count = Column(Integer) # extraction only
     type = Column(String)
     finished = Column(Boolean)
 
diff --git a/software.py b/software.py
index 19fd53e..71ee9e9 100644
--- a/software.py
+++ b/software.py
@@ -96,25 +96,27 @@ class Software:
 
         os.makedirs(software_path, exist_ok=True)
 
+        # beginn progress tracking #
         with zipfile.ZipFile(cache_src, 'r') as zip_ref:
 
+            statekeeper.log_begin_download(local_path=cache_src, path=cache_src, url=None, type="extraction", start_size=len(zip_ref.infolist()))
             total_count = zip_ref.infolist()
             count = 0
             for member in tqdm.tqdm(total_count, desc='Extracting '):
                 try:
                     zip_ref.extract(member, software_path)
                     count += 1
-                    #self.progress_bar_wrapper.get_pb().set(count/len(total_count))
-                    #self.progress_bar_wrapper.get_pb().update_idletasks()
-                    #self.progress_bar_wrapper.set_text(
-                    #                text="Extracting: {:.2f}%".format(count/len(total_count)*100))
+
+                    # update progress #
+                    statekeeper.set_extraction_status(cache_src, count)
+
+
                 except zipfile.error as e:
                     print(e)
                     pass # TODO ???
-            #zip_ref.extractall(software_path)
-
-        #self.progress_bar_wrapper.set_text(text="Loading..")
-        #self.progress_bar_wrapper.update()
+        
+        # finish extraction tracking #
+        statekeeper.log_end_download(cache_src, type="extraction")
 
     def install_async(self):
 
diff --git a/statekeeper.py b/statekeeper.py
index 143e5b9..7e14dd0 100644
--- a/statekeeper.py
+++ b/statekeeper.py
@@ -41,27 +41,45 @@ def _download(url, path):
 
         raise AssertionError("Non-200 Response for:", url, path, response.status_code, response.text)
 
-def log_begin_download(path, local_path, url):
+def log_begin_download(path, local_path, url, type="download", start_size=-1):
 
+    if type == "extraction":
+        print("Extraction path:", path)
+    else:
+        print("Download path", path)
+    
     session = db.session()
-    print("Download path", path)
-    path_exists = session.query(Download).filter(and_(Download.path==path, Download.finished==False)).first()
+    path_exists = session.query(Download).filter(and_(Download.path==path, Download.finished==False, Download.type==type)).first()
+
     if path_exists and False: # TODO FIX THIS
         print("DAFUG", path_exists)
         print("WTF", path_exists.path)
         raise AssertionError("ERROR: {} is already downloading.".format(path))
     else:
         print("Adding to download log:", path)
-        session.merge(Download(path=path, size=-1, type="download", local_path=local_path, url=url, finished=False))
+        session.merge(Download(path=path, size=start_size, type=type, local_path=local_path, url=url, finished=False, count=1))
         session.commit()
 
     db.close_session()
 
-def log_end_download(path):
+def set_extraction_status(path, count):
+
+    session = db.session()
+    obj = session.query(Download).filter(and_(Download.path==path, Download.type=="extraction")).first()
+    if not obj:
+        print("ERROR: {} is not currently extraction, cannot set status.".format(path))
+    else:
+        obj.count = count
+        session.merge(obj)
+        session.commit()
+
+    db.close_session()
+
+def log_end_download(path, type="download"):
 
     print("Downlod end logged", path)
     session = db.session()
-    obj = session.query(Download).filter(Download.path==path).first()
+    obj = session.query(Download).filter(and_(Download.path==path, Download.type==type)).first()
     if not obj:
         raise AssertionError("ERROR: {} is not downloading/cannot remove.".format(path))
     else:
@@ -100,8 +118,16 @@ def get_percent_filled(path):
 
     session = db.session()
     obj = session.query(Download).filter(Download.path==path, Download.finished==False).first()
+    
+    if not obj:
+        return 100
+
+    if obj.type == "extraction":
+        return obj.count / obj.size * 100
+
     if not obj:
         return 100 # means its finished
+
     size = _bytes_to_mb(os.stat(obj.local_path).st_size)
     total_size = get_download_size(obj.path)
     session.close()