From 5d3f4195084bd1f3b184992083b5b9f688d7f3bb Mon Sep 17 00:00:00 2001
From: Leo Vasanko <leo@git>
Date: Mon, 23 Oct 2023 05:24:54 +0300
Subject: [PATCH] /api/watch also continuously reports disk usage

---
 cista/api.py      |  5 +++--
 cista/watching.py | 33 ++++++++++++++++++++++++---------
 2 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/cista/api.py b/cista/api.py
index 9d74738..262e5f9 100644
--- a/cista/api.py
+++ b/cista/api.py
@@ -77,8 +77,9 @@ async def watch(req, ws):
     try:
         with watching.tree_lock:
             q = watching.pubsub[ws] = asyncio.Queue()
-            # Init with full tree
-            await ws.send(watching.refresh())
+            # Init with disk usage and full tree
+            await ws.send(watching.format_du())
+            await ws.send(watching.format_tree())
         # Send updates
         while True:
             await ws.send(await q.get())
diff --git a/cista/watching.py b/cista/watching.py
index 6b76b09..5c1fb0c 100755
--- a/cista/watching.py
+++ b/cista/watching.py
@@ -1,4 +1,5 @@
 import asyncio
+import shutil
 import threading
 import time
 from pathlib import Path, PurePosixPath
@@ -8,7 +9,6 @@ import msgspec
 
 from cista import config
 from cista.protocol import DirEntry, FileEntry, UpdateEntry
-from cista.util.apphelpers import websocket_wrapper
 
 pubsub = {}
 tree = {"": None}
@@ -16,15 +16,18 @@ tree_lock = threading.Lock()
 rootpath = None
 quit = False
 modified_flags = "IN_CREATE", "IN_DELETE", "IN_DELETE_SELF", "IN_MODIFY", "IN_MOVE_SELF", "IN_MOVED_FROM", "IN_MOVED_TO"
+disk_usage = None
 
 def watcher_thread(loop):
+    global disk_usage
+
     while True:
         i = inotify.adapters.InotifyTree(rootpath.as_posix())
-        old = refresh() if tree[""] else None
+        old = format_tree() if tree[""] else None
         with tree_lock:
             # Initialize the tree from filesystem
             tree[""] = walk(rootpath)
-        msg = refresh()
+        msg = format_tree()
         if msg != old:
             asyncio.run_coroutine_threadsafe(broadcast(msg), loop)
 
@@ -33,6 +36,10 @@ def watcher_thread(loop):
 
         for event in i.event_gen():
             if quit: return
+            du = shutil.disk_usage(rootpath)
+            if du != disk_usage:
+                disk_usage = du
+                asyncio.run_coroutine_threadsafe(broadcast(format_du()), loop)
             if time.monotonic() > refreshdl: break
             if event is None: continue
             _, flags, path, filename = event
@@ -43,6 +50,20 @@ def watcher_thread(loop):
             update(path.relative_to(rootpath), loop)
         i = None  # Free the inotify object
 
+def format_du():
+    return msgspec.json.encode({"space": {
+        "disk": disk_usage.total,
+        "used": disk_usage.used,
+        "free": disk_usage.free,
+        "storage": tree[""].size,
+    }}).decode()
+
+def format_tree():
+    root = tree[""]
+    return msgspec.json.encode({"update": [
+        UpdateEntry(size=root.size, mtime=root.mtime, dir=root.dir)
+    ]}).decode()
+
 def walk(path: Path) -> DirEntry | FileEntry | None:
     try:
         s = path.stat()
@@ -63,12 +84,6 @@ def walk(path: Path) -> DirEntry | FileEntry | None:
         print("OS error walking path", path, e)
         return None
 
-def refresh():
-    root = tree[""]
-    return msgspec.json.encode({"update": [
-        UpdateEntry(size=root.size, mtime=root.mtime, dir=root.dir)
-    ]}).decode()
-
 def update(relpath: Path, loop):
     """Called by inotify updates, check the filesystem and broadcast any changes."""
     new = walk(rootpath / relpath)