From 98c0f27633ebf280e786a229fd1fd3532039ae8f Mon Sep 17 00:00:00 2001
From: Rodrigo Braz <rfsbraz@gmail.com>
Date: Sun, 18 Feb 2024 08:38:27 +0000
Subject: [PATCH] feat: Add support for disk size filter (#67)

---
 app/deleterr.py     | 36 ++++++++++++++++++++++++++++++++++--
 app/utils.py        | 15 ++++++++++++++-
 tests/test_utils.py | 14 +++++++++++++-
 3 files changed, 61 insertions(+), 4 deletions(-)

diff --git a/app/deleterr.py b/app/deleterr.py
index 7d6703a..8f53251 100644
--- a/app/deleterr.py
+++ b/app/deleterr.py
@@ -15,7 +15,7 @@
 from app import logger
 from plexapi.server import PlexServer
 from plexapi.exceptions import NotFound
-from app.utils import print_readable_freed_space
+from app.utils import print_readable_freed_space, parse_size_to_bytes
 from app.config import load_config
 from pyarr.exceptions import PyarrResourceNotFound, PyarrServerError
 
@@ -55,6 +55,32 @@ def __init__(self, config):
         self.process_sonarr()
         self.process_radarr()
 
+    def library_meets_disk_space_threshold(self, library, pyarr):
+        for item in library.get("disk_size_threshold"):
+            path = item.get("path")
+            threshold = item.get("threshold")
+            disk_space = pyarr.get_disk_space()
+            folder_found = False
+            for folder in disk_space:
+                if folder["path"] == path:
+                    folder_found = True	
+                    free_space = folder["freeSpace"]
+                    logger.debug(
+                        f"Free space for '{path}': {print_readable_freed_space(free_space)} (threshold: {threshold})"
+                    )
+                    if free_space > parse_size_to_bytes(threshold):
+                        logger.info(
+                            f"Skipping library '{library.get('name')}' as free space is above threshold ({print_readable_freed_space(free_space)} > {threshold})"
+                        )
+                        return False
+            if not folder_found:
+                logger.error(
+                    f"Could not find folder '{path}' in server instance. Skipping library '{library.get('name')}'"
+                )
+                return False
+            
+            return True
+
     def delete_series(self, sonarr, sonarr_show):
         ## PyArr doesn't support deleting the series files, so we need to do it manually
         episodes = sonarr.get_episode(sonarr_show["id"], series=True)
@@ -98,6 +124,9 @@ def process_sonarr(self):
             saved_space = 0
             for library in self.config.settings.get("libraries", []):
                 if library.get("sonarr") == name:
+                    if not self.library_meets_disk_space_threshold(library, sonarr):
+                        continue
+
                     all_show_data = [
                         show
                         for show in unfiltered_all_show_data
@@ -198,6 +227,9 @@ def process_radarr(self):
             saved_space = 0
             for library in self.config.settings.get("libraries", []):
                 if library.get("radarr") == name:
+                    if not self.library_meets_disk_space_threshold(library, radarr):
+                        continue
+
                     max_actions_per_run = _get_config_value(
                         library, "max_actions_per_run", DEFAULT_MAX_ACTIONS_PER_RUN
                     )
@@ -693,7 +725,7 @@ def main():
     logger.info("Running version %s", get_file_contents("/app/commit_tag.txt"))
     logger.info("Log level set to %s", log_level)
 
-    config = load_config("/config/settings.yaml")
+    config = load_config("config/settings.yaml")
     config.validate()
 
     Deleterr(config)
diff --git a/app/utils.py b/app/utils.py
index ee2ede1..3778bdc 100644
--- a/app/utils.py
+++ b/app/utils.py
@@ -6,4 +6,17 @@ def print_readable_freed_space(saved_space):
         saved_space /= 1024
         index += 1
     
-    return f"{saved_space:.2f} {units[index]}"
\ No newline at end of file
+    return f"{saved_space:.2f} {units[index]}"
+
+def parse_size_to_bytes(size_str):
+    units = ["B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"]
+    unit = ''.join([i for i in size_str if not i.isdigit() and not i == '.'])
+    size = ''.join([i for i in size_str if i.isdigit() or i == '.'])
+    size = float(size)
+    index = units.index(unit)
+    
+    while index > 0:
+        size *= 1024
+        index -= 1
+    
+    return int(size)
\ No newline at end of file
diff --git a/tests/test_utils.py b/tests/test_utils.py
index e99ba0c..bcb4ee3 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -1,5 +1,5 @@
 import pytest
-from app.utils import print_readable_freed_space
+from app.utils import print_readable_freed_space, parse_size_to_bytes
 
 @pytest.mark.parametrize("input_size, expected_output", [
     (500, "500.00 Bytes"),
@@ -12,3 +12,15 @@
 def test_print_readable_freed_space(input_size, expected_output):
     result = print_readable_freed_space(input_size)
     assert result == expected_output, f"For {input_size}, expected {expected_output} but got {result}"
+
+@pytest.mark.parametrize("input_size, expected_output", [
+    ('2TB', 2199023255552),
+    ('5.6GB', 6012954214),
+    ("1GB", 1073741824),
+    ('230MB', 241172480), 
+    ('1B', 1),
+    ('0KB', 0),
+])
+def test_print_readable_freed_space(input_size, expected_output):
+    result = parse_size_to_bytes(input_size)
+    assert result == expected_output, f"For {input_size}, expected {expected_output} but got {result}"
\ No newline at end of file