Add types and clean up

download_file.py

@@ -1,7 +1,7 @@
 import requests
 
 
-async def download_file(url, destination):
+async def download_file(url: str, destination: str) -> None:
     try:
         response = requests.get(url, stream=True)
         response.raise_for_status()  # Check for HTTP errors
@@ -12,5 +12,12 @@ async def download_file(url, destination):
                     f.write(chunk)
 
         print(f"File downloaded to: {destination}")
-    except requests.exceptions.RequestException as e:
+
+    except requests.exceptions.HTTPError as e:
+        print(
+            f"Download failed: HTTP Error - {e.response.status_code} {e.response.reason}"
+        )
+    except requests.exceptions.Timeout as e:
+        print(f"Download failed: Request timed out.")
+    except requests.exceptions.RequestException as e:  # Fallback for other errors
         print(f"Download failed: {e}")

updater.py

@@ -4,13 +4,17 @@
 from zipfile import ZipFile
 from pymongo import MongoClient
 from download_file import download_file
+from utils.create_unique_key import create_unique_key
+from typing import List, Dict, Any
 
 EXTRACT_PATH = "tmp"
 
-csv_data = []
+csv_data: List[Dict[str, Any]] = []
 
 
-async def update(collection_name, zip_file_name, zip_url, key_fields):
+async def update(
+    collection_name: str, zip_file_name: str, zip_url: str, key_fields: List[str]
+) -> str:
     client = MongoClient(os.environ.get("MONGODB_URI", "mongodb://localhost:27017/"))
     db = client[os.environ.get("MONGODB_DB_NAME", "local-lta-datasets")]
     collection = db[collection_name]
@@ -30,9 +34,6 @@ async def update(collection_name, zip_file_name, zip_url, key_fields):
 
         existing_data = collection.find({})
 
-        def create_unique_key(item, key_fields):
-            return "-".join(str(item[field]) for field in key_fields if item.get(field))
-
         existing_data_map = {
             create_unique_key(item, key_fields): item for item in existing_data
         }
@@ -52,20 +53,22 @@ def create_unique_key(item, key_fields):
 
         return message
 
-    except Exception as error:  # Use Exception for a broader error catch
+    except Exception as error:
         print(f"An error has occurred: {error}")
         raise
 
 
-def extract_zip_file(zip_file_path, extract_to_path):
+def extract_zip_file(zip_file_path: str, extract_to_path: str) -> str:
     with ZipFile(zip_file_path, "r") as zip_ref:
         zip_ref.extractall(extract_to_path)
         for entry in zip_ref.infolist():
             if not entry.is_dir():
                 return entry.filename
 
 
-async def main(collection_name, zip_file_name, zip_url, key_fields):
+async def main(
+    collection_name: str, zip_file_name: str, zip_url: str, key_fields: List[str]
+) -> Dict[str, Any]:
     message = await update(
         collection_name=collection_name,
         zip_file_name=zip_file_name,

utils/create_unique_key.py

@@ -0,0 +1,5 @@
+from typing import Dict, Any, List
+
+
+def create_unique_key(item: Dict[str, Any], key_fields: List[str]) -> str:
+    return "-".join(str(item[field]) for field in key_fields if item.get(field))