Extract input zip files before processing them (#6)

my_proof/__main__.py

@@ -3,6 +3,7 @@
 import os
 import sys
 import traceback
+import zipfile
 from typing import Dict, Any
 
 from my_proof.proof import Proof
@@ -31,6 +32,7 @@ def run() -> None:
 
     if not input_files_exist:
         raise FileNotFoundError(f"No input files found in {INPUT_DIR}")
+    extract_input()
 
     proof = Proof(config)
     proof_response = proof.generate()
@@ -41,6 +43,19 @@ def run() -> None:
     logging.info(f"Proof generation complete: {proof_response}")
 
 
+def extract_input() -> None:
+    """
+    If the input directory contains any zip files, extract them
+    :return:
+    """
+    for input_filename in os.listdir(INPUT_DIR):
+        input_file = os.path.join(INPUT_DIR, input_filename)
+
+        if zipfile.is_zipfile(input_file):
+            with zipfile.ZipFile(input_file, 'r') as zip_ref:
+                zip_ref.extractall(INPUT_DIR)
+
+
 if __name__ == "__main__":
     try:
         run()

my_proof/proof.py

@@ -20,18 +20,20 @@ def generate(self) -> ProofResponse:
         # Iterate through files and calculate data validity
         account_email = None
         total_score = 0
+
         for input_filename in os.listdir(self.config['input_dir']):
             input_file = os.path.join(self.config['input_dir'], input_filename)
-            with open(input_file, 'r') as f:
-                input_data = json.load(f)
+            if os.path.splitext(input_file)[1].lower() == '.json':
+                with open(input_file, 'r') as f:
+                    input_data = json.load(f)
 
-                if input_filename == 'account.json':
-                    account_email = input_data.get('email', None)
-                    continue
+                    if input_filename == 'account.json':
+                        account_email = input_data.get('email', None)
+                        continue
 
-                elif input_filename == 'activity.json':
-                    total_score = sum(item['score'] for item in input_data)
-                    continue
+                    elif input_filename == 'activity.json':
+                        total_score = sum(item['score'] for item in input_data)
+                        continue
 
         email_matches = self.config['user_email'] == account_email
         score_threshold = fetch_random_number()