refactor Python checks for improved validation

custom_functions/prechecks/cdm_checks.py

@@ -1,88 +1,106 @@
+"""Module to check for and download CDM files."""
+
 import os
+import sys
 import yaml
 import requests
 
+CONFIG_PATH = os.path.join(os.getcwd(), "configs", "config.yaml")
 
-def check_for_wvd_cdm():
-    with open(f"{os.getcwd()}/configs/config.yaml", "r") as file:
-        config = yaml.safe_load(file)
-    if config["default_wv_cdm"] == "":
-        answer = " "
-        while answer[0].upper() != "Y" and answer[0].upper() != "N":
-            answer = input(
-                "No default Widevine CDM specified, would you like to download one from The CDM Project? (Y)es/(N)o: "
-            )
-        if answer[0].upper() == "Y":
-            response = requests.get(
-                url="https://cdm-project.com/CDRM-Team/CDMs/raw/branch/main/Widevine/L3/public.wvd"
-            )
+
+def load_config():
+    """Load the config file."""
+    with open(CONFIG_PATH, "r", encoding="utf-8") as file:
+        return yaml.safe_load(file)
+
+
+def save_config(config):
+    """Save the config file."""
+    with open(CONFIG_PATH, "w", encoding="utf-8") as file:
+        yaml.dump(config, file)
+
+
+def prompt_yes_no(message):
+    """Prompt the user for a yes or no answer."""
+    answer = " "
+    while answer[0].upper() not in ["Y", "N"]:
+        answer = input(message)
+    return answer[0].upper() == "Y"
+
+
+def check_for_cdm(config_key, file_ext, download_url, cdm_dir, cdm_name):
+    """Check for a CDM file."""
+    config = load_config()
+    cdm_value = config.get(config_key, "")
+    cdm_dir_path = os.path.join(os.getcwd(), "configs", "CDMs", cdm_dir)
+    os.makedirs(cdm_dir_path, exist_ok=True)
+
+    if not cdm_value:
+        if prompt_yes_no(
+            f"No default {cdm_name} CDM specified, would you like to download one "
+            "from The CDM Project? (Y)es / (N)o: "
+        ):
+            response = requests.get(download_url, timeout=10)
             if response.status_code == 200:
-                with open(f"{os.getcwd()}/configs/CDMs/WV/public.wvd", "wb") as file:
+                file_path = os.path.join(cdm_dir_path, f"public.{file_ext}")
+                with open(file_path, "wb") as file:
                     file.write(response.content)
-                config["default_wv_cdm"] = "public"
-                with open(f"{os.getcwd()}/configs/config.yaml", "w") as file:
-                    yaml.dump(config, file)
-                print("Successfully downloaded Widevine CDM")
+                config[config_key] = "public"
+                save_config(config)
+                print(f"Successfully downloaded {cdm_name} CDM")
             else:
-                exit(
-                    f"Download failed, please try again or place a .wvd file in {os.getcwd()}/configs/CDMs/WV and specify the name in {os.getcwd()}/configs/config.yaml"
+                sys.exit(
+                    f"Download failed, please try again, or place a .{file_ext} file "
+                    f"in {cdm_dir_path} and specify the name in {CONFIG_PATH}"
                 )
-        if answer[0].upper() == "N":
-            exit(
-                f"Place a .wvd file in {os.getcwd()}/configs/CDMs/WV and specify the name in {os.getcwd()}/configs/config.yaml"
+        else:
+            sys.exit(
+                f"Place a .{file_ext} file in {cdm_dir_path} and specify the name in {CONFIG_PATH}"
             )
     else:
-        base_name = config["default_wv_cdm"]
-        if not base_name.endswith(".wvd"):
-            base_name += ".wvd"
-        if os.path.exists(f"{os.getcwd()}/configs/CDMs/WV/{base_name}"):
+        base_name = (
+            cdm_value
+            if cdm_value.endswith(f".{file_ext}")
+            else f"{cdm_value}.{file_ext}"
+        )
+        file_path = os.path.join(cdm_dir_path, base_name)
+        if os.path.exists(file_path):
             return
-        else:
-            exit(
-                f"Widevine CDM {base_name} does not exist in {os.getcwd()}/configs/CDMs/WV"
-            )
-
-
-def check_for_prd_cdm():
-    with open(f"{os.getcwd()}/configs/config.yaml", "r") as file:
-        config = yaml.safe_load(file)
-    if config["default_pr_cdm"] == "":
-        answer = " "
-        while answer[0].upper() != "Y" and answer[0].upper() != "N":
-            answer = input(
-                "No default PlayReady CDM specified, would you like to download one from The CDM Project? (Y)es/(N)o: "
-            )
-        if answer[0].upper() == "Y":
-            response = requests.get(
-                url="https://cdm-project.com/CDRM-Team/CDMs/raw/branch/main/Playready/SL2000/public.prd"
-            )
+        # Prompt to download if file is missing, even if config has a value
+        if prompt_yes_no(
+            f"{cdm_name} CDM {base_name} does not exist in {cdm_dir_path}. Would you like to download it from The CDM Project? (Y)es/(N)o: "
+        ):
+            response = requests.get(download_url, timeout=10)
             if response.status_code == 200:
-                with open(f"{os.getcwd()}/configs/CDMs/PR/public.prd", "wb") as file:
+                with open(file_path, "wb") as file:
                     file.write(response.content)
-                config["default_pr_cdm"] = "public"
-                with open(f"{os.getcwd()}/configs/config.yaml", "w") as file:
-                    yaml.dump(config, file)
-                print("Successfully downloaded PlayReady CDM")
+                config[config_key] = base_name.replace(f".{file_ext}", "")
+                save_config(config)
+                print(f"Successfully downloaded {cdm_name} CDM")
             else:
-                exit(
-                    f"Download failed, please try again or place a .prd file in {os.getcwd()}/configs/CDMs/PR and specify the name in {os.getcwd()}/configs/config.yaml"
+                sys.exit(
+                    f"Download failed, please try again, or place a .{file_ext} file "
+                    f"in {cdm_dir_path} and specify the name in {CONFIG_PATH}"
                 )
-        if answer[0].upper() == "N":
-            exit(
-                f"Place a .prd file in {os.getcwd()}/configs/CDMs/PR and specify the name in {os.getcwd()}/configs/config.yaml"
-            )
-    else:
-        base_name = config["default_pr_cdm"]
-        if not base_name.endswith(".prd"):
-            base_name += ".prd"
-        if os.path.exists(f"{os.getcwd()}/configs/CDMs/PR/{base_name}"):
-            return
         else:
-            exit(
-                f"PlayReady CDM {base_name} does not exist in {os.getcwd()}/configs/CDMs/WV"
+            sys.exit(
+                f"Place a .{file_ext} file in {cdm_dir_path} and specify the name in {CONFIG_PATH}"
             )
 
 
 def check_for_cdms():
-    check_for_wvd_cdm()
-    check_for_prd_cdm()
+    """Check for CDM files."""
+    check_for_cdm(
+        config_key="default_wv_cdm",
+        file_ext="wvd",
+        download_url="https://cdm-project.com/CDRM-Team/CDMs/raw/branch/main/Widevine/L3/public.wvd",
+        cdm_dir="WV",
+        cdm_name="Widevine",
+    )
+    check_for_cdm(
+        config_key="default_pr_cdm",
+        file_ext="prd",
+        download_url="https://cdm-project.com/CDRM-Team/CDMs/raw/branch/main/Playready/SL2000/public.prd",
+        cdm_dir="PR",
+        cdm_name="PlayReady",
+    )

custom_functions/prechecks/python_checks.py

@@ -1,33 +1,26 @@
+"""Module to check for the Python version and environment."""
+
 import sys
 import os
 import subprocess
 import venv
+import importlib.util
 
 
 def version_check():
-    major_version = sys.version_info.major
-    minor_version = sys.version_info.minor
-
-    if major_version >= 3:
-        if minor_version >= 12:
-            return
-        else:
-            exit("Python version 3.12 or higher is required")
-    else:
-        exit("Python 2 detected, Python version 3.12 or higher is required")
+    """Check for the Python version."""
+    if sys.version_info < (3, 12):
+        sys.exit("Python version 3.12 or higher is required")
 
 
 def pip_check():
-    try:
-        import pip
-
-        return
-    except ImportError:
-        exit("Pip is not installed")
+    """Check for the pip installation."""
+    if importlib.util.find_spec("pip") is None:
+        sys.exit("Pip is not installed")
 
 
 def venv_check():
-    # Check if we're already inside a virtual environment
+    """Check for the virtual environment."""
     if hasattr(sys, "real_prefix") or (
         hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix
     ):
@@ -40,20 +33,20 @@ def venv_check():
         else os.path.join(venv_path, "Scripts", "python.exe")
     )
 
-    # If venv already exists, restart script using its Python
     if os.path.exists(venv_path):
         subprocess.call([venv_python] + sys.argv)
         sys.exit()
 
-    # Ask user for permission to create a virtual environment
-    answer = ""
-    while not answer or answer[0].upper() not in {"Y", "N"}:
-        answer = input(
-            "Program is not running from a venv. To maintain compatibility and dependencies, this program must be run from one.\n"
-            "Would you like me to create one for you? (Y/N): "
+    answer = (
+        input(
+            "Program is not running from a virtual environment. To maintain "
+            "compatibility, this program must be run from one.\n"
+            "Would you like to create one? (Y/N): "
         )
-
-    if answer[0].upper() == "Y":
+        .strip()
+        .upper()
+    )
+    if answer.startswith("Y"):
         print("Creating virtual environment...")
         venv.create(venv_path, with_pip=True)
         subprocess.call([venv_python] + sys.argv)
@@ -64,37 +57,54 @@ def venv_check():
 
 
 def requirements_check():
-    try:
-        import pywidevine
-        import pyplayready
-        import flask
-        import flask_cors
-        import yaml
-        import mysql.connector
-
+    """Check for the requirements."""
+    required_packages = [
+        "pywidevine",
+        "pyplayready",
+        "flask",
+        "flask_cors",
+        "yaml",
+        "mysql.connector",
+    ]
+    missing = []
+    for pkg in required_packages:
+        if "." in pkg:
+            parent, _ = pkg.split(".", 1)
+            if (
+                importlib.util.find_spec(parent) is None
+                or importlib.util.find_spec(pkg) is None
+            ):
+                missing.append(pkg)
+        else:
+            if importlib.util.find_spec(pkg) is None:
+                missing.append(pkg)
+    if not missing:
         return
-    except ImportError:
-        while True:
-            user_input = (
-                input("Missing packages. Do you want to install them? (Y/N): ")
-                .strip()
-                .upper()
+
+    while True:
+        user_input = (
+            input(
+                f"Missing packages: {', '.join(missing)}. Do you want to install them? (Y/N): "
             )
-            if user_input == "Y":
-                print("Installing packages from requirements.txt...")
-                subprocess.check_call(
-                    [sys.executable, "-m", "pip", "install", "-r", "requirements.txt"]
-                )
-                print("Installation complete.")
-                break
-            elif user_input == "N":
-                print("Dependencies required, please install them and run again.")
-                sys.exit()
-            else:
-                print("Invalid input. Please enter 'Y' to install or 'N' to exit.")
+            .strip()
+            .upper()
+        )
+        if user_input == "Y":
+            print("Installing packages from requirements.txt...")
+            subprocess.check_call(
+                [sys.executable, "-m", "pip", "install", "-r", "requirements.txt"]
+            )
+            print("Installation complete.")
+            break
+        if user_input == "N":
+            print("Dependencies required, please install them and run again.")
+            sys.exit()
+        else:
+            print("Invalid input. Please enter 'Y' to install or 'N' to exit.")
 
 
 def run_python_checks():
+    """Run the Python checks."""
     if getattr(sys, "frozen", False):  # Check if running from PyInstaller
         return
     version_check()