fix: useless repeated code

g0ldyy · Jul 11, 2024 · fa8ce05 · fa8ce05
1 parent db1a003
commit fa8ce05
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 26 deletions.
diff --git a/comet/api/core.py b/comet/api/core.py
@@ -5,7 +5,7 @@
 from fastapi.templating import Jinja2Templates
 
 from comet.utils.models import settings
-from comet.utils.general import config_check
+from comet.utils.general import config_check, get_debrid_extension
 
 templates = Jinja2Templates("comet/templates")
 main = APIRouter()
@@ -65,15 +65,7 @@ async def manifest(b64config: str = None):
     if not config:
         config = {"debridService": None}
 
-    debrid_extension = None
-    if config["debridService"] == "realdebrid":
-        debrid_extension = "RD"
-    elif config["debridService"] == "alldebrid":
-        debrid_extension = "AD"
-    elif config["debridService"] == "premiumize":
-        debrid_extension = "PM"
-    elif config["debridService"] == "torbox":
-        debrid_extension = "TB"
+    debrid_extension = get_debrid_extension(config["debridService"])
 
     return {
         "id": settings.ADDON_ID,

diff --git a/comet/api/stream.py b/comet/api/stream.py
@@ -14,6 +14,7 @@
 from comet.utils.general import (
     bytes_to_size,
     config_check,
+    get_debrid_extension,
     get_indexer_manager,
     get_zilean,
     filter,
@@ -128,14 +129,7 @@ async def stream(request: Request, b64config: str, type: str, id: str):
                 )
                 sorted_ranked_files = json.loads(sorted_ranked_files[0])
 
-                if config["debridService"] == "realdebrid":
-                    debrid_extension = "RD"
-                elif config["debridService"] == "alldebrid":
-                    debrid_extension = "AD"
-                elif config["debridService"] == "premiumize":
-                    debrid_extension = "PM"
-                elif config["debridService"] == "torbox":
-                    debrid_extension = "TB"
+                debrid_extension = get_debrid_extension(config["debridService"])
 
                 balanced_hashes = await get_balanced_hashes(sorted_ranked_files, config)
 
@@ -278,14 +272,7 @@ async def stream(request: Request, b64config: str, type: str, id: str):
         )
         logger.info(f"Results have been cached for {log_name}")
 
-        if config["debridService"] == "realdebrid":
-            debrid_extension = "RD"
-        elif config["debridService"] == "alldebrid":
-            debrid_extension = "AD"
-        elif config["debridService"] == "premiumize":
-            debrid_extension = "PM"
-        elif config["debridService"] == "torbox":
-            debrid_extension = "TB"
+        debrid_extension = get_debrid_extension(config["debridService"])
 
         balanced_hashes = await get_balanced_hashes(sorted_ranked_files, config)
 

diff --git a/comet/utils/general.py b/comet/utils/general.py
@@ -190,6 +190,20 @@ def config_check(b64config: str):
         return False
 
 
+def get_debrid_extension(debridService: str):
+    debrid_extension = "?"  # Unknown
+    if debridService == "realdebrid":
+        debrid_extension = "RD"
+    elif debridService == "alldebrid":
+        debrid_extension = "AD"
+    elif debridService == "premiumize":
+        debrid_extension = "PM"
+    elif debridService == "torbox":
+        debrid_extension = "TB"
+
+    return debrid_extension
+
+
 async def get_indexer_manager(
     session: aiohttp.ClientSession,
     indexer_manager_type: str,