From 26a94aa02dd1237798788dfd36795e3e248772e8 Mon Sep 17 00:00:00 2001
From: John Lancaster <32917998+jsl12@users.noreply.github.com>
Date: Mon, 27 May 2024 12:49:36 -0500
Subject: [PATCH] fleshed out prune module

---
 src/restic/prune.py | 53 +++++++++++++++++++++++++++++++++++----------
 1 file changed, 42 insertions(+), 11 deletions(-)

diff --git a/src/restic/prune.py b/src/restic/prune.py
index 3d9128e..7b5e71c 100644
--- a/src/restic/prune.py
+++ b/src/restic/prune.py
@@ -1,33 +1,64 @@
+import json
 import logging
 import re
 import subprocess
+import sys
+
+from rich.logging import RichHandler
 
 from restic.console import console, logger
+from restic.loki import send_to_loki
+
+field_regex = re.compile(
+    r'^(?P<name>[\w ]+):\s+(?P<blobs>\d+) blobs \/ (?P<size>\d+(\.\d+)? (M|G)iB)', re.MULTILINE
+)
 
 
-def main(dry_run: bool = False):
+def convert_size(size_str: str) -> int:
+    base_size = float(size_str.split(' ')[0])
+
+    if size_str.endswith('GiB'):
+        scale = 8589934592
+    elif size_str.endswith('MiB'):
+        scale = 8388608
+    return int(round(base_size * scale))
+
+
+def field_gen(result_str: str):
+    for m in field_regex.finditer(result_str):
+        d = m.groupdict()
+        d['blobs'] = int(d['blobs'])
+        d['size'] = convert_size(d['size'])
+        yield d
+
+
+def prune(loki_url: str = None, dry_run: bool = False):
     cmd = ['restic', 'prune']
     if dry_run:
-        cmd.append('-n')
+        cmd.append('--dry-run')
     logger.debug(f'Running cmd [bright_black]{" ".join(cmd)}[/]')
 
     with console.status('Pruning...'):
         result = subprocess.run(cmd, capture_output=True, text=True)
 
-    if m := re.search(r'total prune:\s+(.*?)$', result.stdout, re.MULTILINE):
-        pruned_data = m.group(1).split(' / ')[1]
-        logger.info(f'Pruned {pruned_data}')
+    if result.returncode != 0:
+        logger.error(result.stderr)
+        sys.exit(1)
 
-    if m := re.search(r'remaining:\s+(.*?)$', result.stdout, re.MULTILINE):
-        remaining = m.group(1).split(' / ')[1]
-        logger.info(f'{remaining} remaining')
+    d = {f['name']: {'blobs': f['blobs'], 'size': f['size']} for f in field_gen(result.stdout)}
+    console.print(d)
+
+    if loki_url is not None and not dry_run:
+        send_to_loki(loki_url, line=json.dumps(d), backup='prune')
 
 
-if __name__ == '__main__':
-    from rich.logging import RichHandler
-
+def main():
     logging.basicConfig(
         level='DEBUG', format='%(message)s', handlers=[RichHandler(markup=True, console=console)]
     )
 
+    prune(dry_run=True)
+
+
+if __name__ == '__main__':
     main()