feat: cleanup

buroa · Oct 8, 2024 · 9959ff1 · 9959ff1
1 parent 512cc46
commit 9959ff1
Showing 1 changed file with 100 additions and 43 deletions.
diff --git a/qbtools/commands/limiter.py b/qbtools/commands/limiter.py
@@ -1,36 +1,112 @@
 import time
 import requests
-import logging
-from decimal import Decimal, InvalidOperation
+
+from qbtools import utils
+from decimal import Decimal
 from typing import Optional, Tuple
 
+
 def __init__(app, logger):
     logger.info("Starting limiter process...")
 
     while True:
-        qbittorrent_queue, qbittorrent_current_limit = qbittorrent_data(app)
-        sabnzbd_queue, sabnzbd_current_limit = sabnzbd_data(app)
-
-        logger.debug(f"qBittorrent: {qbittorrent_queue} item(s) @ {qbittorrent_current_limit / 1024 /1024 } MB/s")
-        logger.debug(f"SabNZBD: {sabnzbd_queue} item(s) @ max {sabnzbd_current_limit} MB/s")
-
-        percentage = app.limit_percent if qbittorrent_queue > 0 and sabnzbd_queue > 0 else app.max_percent
-
-        qbittorrent_updated_limit = int(app.max_line_speed_mbps * percentage * 1024 * 1024)
-
-        if qbittorrent_current_limit != qbittorrent_updated_limit:
-            app.client.transfer_set_download_limit(qbittorrent_updated_limit)
-            logger.info(f"qbittorrent download limit set to {qbittorrent_updated_limit / 1024 / 1024 } MB/s (was {qbittorrent_current_limit} MB/s)...")
+        qbittorrent_queue, qbittorrent_current_limit = qbittorrent_data(app, logger)
+        sabnzbd_queue, sabnzbd_current_limit = sabnzbd_data(app, logger)
+
+        logger.debug(
+            f"qBittorrent: {qbittorrent_queue} item(s) @ {qbittorrent_current_limit} MB/s"
+        )
+        logger.debug(
+            f"SabNZBD: {sabnzbd_queue} item(s) @ max {sabnzbd_current_limit} MB/s"
+        )
+
+        percentage = (
+            app.limit_percent
+            if qbittorrent_queue and sabnzbd_queue
+            else app.max_percent
+        )
+
+        qbittorrent_updated_limit = int(app.max_line_speed_mbps * percentage)
+
+        if (
+            qbittorrent_current_limit is not None
+            and qbittorrent_current_limit != qbittorrent_updated_limit
+        ):
+            try:
+                app.client.transfer_set_download_limit(
+                    qbittorrent_updated_limit * 1024 * 1024
+                )
+            except Exception as e:
+                logger.error(e)
+            else:
+                logger.info(
+                    f"qbittorrent download limit set to "
+                    f"{qbittorrent_updated_limit} MB/s "
+                    f"(was {qbittorrent_current_limit} MB/s)..."
+                )
 
         sabnzbd_updated_limit = int(app.max_line_speed_mbps * percentage)
 
-        if sabnzbd_current_limit != sabnzbd_updated_limit:
-            handle_request(f'{app.sabnzbd_host}/api', 'POST', {'apikey': app.sabnzbd_apikey, 'mode': 'config', 'name': 'speedlimit', 'value': sabnzbd_updated_limit})
-            logger.info(f"sabnzbd download limit set to {sabnzbd_updated_limit} MB/s (was {sabnzbd_current_limit} MB/s)...")
+        if (
+            sabnzbd_current_limit is not None
+            and sabnzbd_current_limit != sabnzbd_updated_limit
+        ):
+            try:
+                handle_request(
+                    url=f"{app.sabnzbd_host}/api",
+                    method="POST",
+                    data=dict(
+                        apikey=app.sabnzbd_apikey,
+                        mode="config",
+                        name="speedlimit",
+                        value=sabnzbd_updated_limit,
+                    ),
+                )
+            except Exception as e:
+                logger.error(e)
+            else:
+                logger.info(
+                    f"sabnzbd download limit set to {sabnzbd_updated_limit} MB/s "
+                    f"(was {sabnzbd_current_limit} MB/s)..."
+                )
 
         time.sleep(app.interval)
 
 
+def qbittorrent_data(app, logger) -> Tuple[int, int]:
+    try:
+        torrents = len(app.client.torrents.info(status_filter="downloading"))
+        download_limit = app.client.transfer_download_limit() / 1024 / 1024
+    except Exception as e:
+        logger.exception(e)
+        return None, None
+    else:
+        return torrents, download_limit
+
+
+def sabnzbd_data(app, logger) -> Tuple[int, int]:
+    try:
+        data = handle_request(
+            f"{app.sabnzbd_host}/api?apikey={app.sabnzbd_apikey}&mode=queue&output=json"
+        )
+    except Exception as e:
+        logger.exception(e)
+        return None, None
+    else:
+        return (
+            int(data.get("queue", {}).get("noofslots", 0)) if data else 0,
+            int(data.get("queue", {}).get("speedlimit", 0)) if data else 0,
+        )
+
+
+def handle_request(
+    url: str, method: str = "GET", data: Optional[dict] = None
+) -> Optional[dict]:
+    response = requests.post(url, data=data) if method == "POST" else requests.get(url)
+    response.raise_for_status()
+    return response.json() if method == "GET" else None
+
+
 def add_arguments(command, subparser):
     """
     Description:
@@ -41,15 +117,17 @@ def add_arguments(command, subparser):
     parser = subparser.add_parser(command)
     parser.add_argument(
         "--sabnzbd-host",
-        type=str,
-        required=True,
+        action=utils.EnvDefault,
+        envvar="SABNZBD_HOST",
         help="The host of the SabNZBD server",
+        required=True,
     )
     parser.add_argument(
         "--sabnzbd-apikey",
-        type=str,
-        required=True,
+        action=utils.EnvDefault,
+        envvar="SABNZBD_API_KEY",
         help="The API key for the SabNZBD server",
+        required=True,
     )
     parser.add_argument(
         "--max-line-speed-mbps",
@@ -75,24 +153,3 @@ def add_arguments(command, subparser):
         default=5,
         help="The interval to check the speeds in seconds",
     )
-
-
-def qbittorrent_data(app) -> Tuple[int, int]:
-    return (
-        len(app.client.torrents.info(status_filter='downloading')),
-        app.client.transfer_download_limit()
-    )
-
-
-def sabnzbd_data(app) -> Tuple[int, int]:
-    data = handle_request(f'{app.sabnzbd_host}/api?apikey={app.sabnzbd_apikey}&mode=queue&output=json')
-    return (
-        int(data.get('queue', {}).get('noofslots', 0)) if data else 0,
-        int(data.get('queue', {}).get('speedlimit', 0)) if data else 0
-    )
-
-
-def handle_request(url: str, method: str = 'GET', data: Optional[dict] = None) -> Optional[dict]:
-    response = requests.post(url, data=data) if method == 'POST' else requests.get(url)
-    response.raise_for_status()
-    return response.json() if method == 'GET' else None