From 557d1d5c19ab518fa7abb3229c6d9042867e6c00 Mon Sep 17 00:00:00 2001
From: Frederick Muriuki Muriithi
Date: Fri, 10 Jun 2022 08:06:47 +0300
Subject: Enable upload of zipfiles

---
 scripts/worker.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

(limited to 'scripts')

diff --git a/scripts/worker.py b/scripts/worker.py
index 0ef5ae5..ecdfaa2 100644
--- a/scripts/worker.py
+++ b/scripts/worker.py
@@ -1,6 +1,7 @@
 import os
 import sys
 from typing import Callable
+from zipfile import Path, ZipFile, is_zipfile
 
 import jsonpickle
 from redis import Redis
@@ -64,6 +65,10 @@ def make_user_aborted(redis_conn, job_id):
         return user_aborted
     return __aborted__
 
+def get_zipfile_size(filepath):
+    with ZipFile(filepath, "r") as zfile:
+        return zfile.infolist()[0].file_size
+
 def main():
     args = process_cli_arguments()
     if args is None:
@@ -72,7 +77,8 @@ def main():
 
     with Redis.from_url(args.redisurl) as redis_conn:
         progress_calculator = make_progress_calculator(
-            os.stat(args.filepath).st_size)
+            get_zipfile_size(args.filepath) if is_zipfile(args.filepath)
+            else os.stat(args.filepath).st_size)
         progress_indicator = make_progress_indicator(
             redis_conn, args.job_id, progress_calculator)
         count = args.count
-- 
cgit v1.2.3