Improve readability and maintainability in downloader.py

supervision/assets/downloader.py

@@ -1,3 +1,4 @@
+import os
 from functools import partial
 from hashlib import md5
 from os import remove as os_remove
@@ -22,31 +23,33 @@ except ImportError:
     )
 
 
-# md5 check function
-def _check_md5(filename: str, orig_md5: str) -> bool:
+def is_md5_hash_matching(filename: str, original_md5_hash: str) -> bool:
     """
-    filename: str, A string representing the path to the file.
-    orig_md5: str, A string representing the original md5 hash.
+    Check if the MD5 hash of a file matches the original hash.
+
+    Parameters:
+        filename (str): The path to the file to be checked as a string.
+        original_md5_hash (str): The original MD5 hash to compare against.
+
+    Returns:
+        bool: True if the hashes match, False otherwise.
     """
-    if not path_exists(filename):
-        return False
-    with open(filename, "rb") as file_to_check:
-        # read contents of the file
-        data = file_to_check.read()
-        # pipe contents of the file through
-        md5_returned = md5(data).hexdigest()
-        # Return True if the computed hash matches the original one
-        if md5_returned == orig_md5:
-            return True
+    if not os.path.exists(filename):
         return False
 
+    with open(filename, "rb") as file:
+        file_contents = file.read()
+        computed_md5_hash = md5(file_contents).hexdigest()
+
+    return computed_md5_hash == original_md5_hash
+
 
 def download_assets(asset_name: Union[VideoAssets, str]) -> str:
     """
     asset_name: VIDEO_ASSETS,  name of the file to download
     """
 
-    filename = f"{asset_name}"
+    filename = asset_name.value if isinstance(asset_name, VideoAssets) else asset_name
     if not path_exists(filename) and asset_name in VIDEO_ASSETS:
         print(f"Downloading {filename} assets \n")
         res: Response = get(
@@ -66,19 +69,19 @@ def download_assets(asset_name: Union[VideoAssets, str]) -> str:
         desc = "(Unknown total file size)" if file_size == 0 else ""
         res.raw.read = partial(
             res.raw.read, decode_content=True
-        )  # Decompress if needed
+        )
         with tqdm.wrapattr(
             res.raw,
             "read",
             total=file_size,
             desc=desc,
-            colour="a351fb",
+            colour="#a351fb",
         ) as r_raw:
             with folder_path.open("wb") as f:
                 copyfileobj(r_raw, f)
 
     elif path_exists(filename):
-        if not _check_md5(filename, VIDEO_ASSETS[asset_name][2]):
+        if not is_md5_hash_matching(filename, VIDEO_ASSETS[asset_name][2]):
             print("File corrupted. Re-downloading... \n")
             os_remove(filename)
             download_assets(asset_name)
@@ -86,14 +89,7 @@ def download_assets(asset_name: Union[VideoAssets, str]) -> str:
         print(f"{filename} asset download complete. \n")
     else:
         raise ValueError(
-            "Invalid asset type. It should be one of the following: \n"
-            "vehicles.mp4 \n"
-            "vehicles-2.mp4 \n"
-            "milk-video-1.mp4 \n"
-            "grocery-store.mp4 \n"
-            "subway.mp4 \n"
-            "market-square.mp4 \n"
-            "people-walking-bw.mp4 \n"
+            f"Invalid asset. It should be one of the following: {VideoAssets.list()}."
         )
 
     return filename