removed _StorageState

2 files changed, 67 insertions, 91 deletions

diff --git a/kvmd/plugins/msd/otg/__init__.py b/kvmd/plugins/msd/otg/__init__.py
index 50b8aa72..d8175975 100644
--- a/kvmd/plugins/msd/otg/__init__.py
+++ b/kvmd/plugins/msd/otg/__init__.py
@@ -67,14 +67,6 @@ class _DriveState:
     rw: bool
 
 
[email protected](frozen=True)
-class _StorageState:
-    size: int
-    free: int
-    images: dict[str, Image]
-
-
-# =====
 @dataclasses.dataclass
 class _VirtualDriveState:
     image: (Image | None)
@@ -96,7 +88,7 @@ class _State:
     def __init__(self, notifier: aiotools.AioNotifier) -> None:
         self.__notifier = notifier
 
-        self.storage: (_StorageState | None) = None
+        self.storage: (Storage | None) = None
         self.vd: (_VirtualDriveState | None) = None
 
         self._region = aiotools.AioExclusiveRegion(MsdIsBusyError)
@@ -175,21 +167,17 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
         async with self.__state._lock:  # pylint: disable=protected-access
             storage: (dict | None) = None
             if self.__state.storage:
+                if self.__writer:
+                    # При загрузке файла показываем актуальную статистику вручную
+                    await self.__storage.reload_size_only()
+
                 storage = dataclasses.asdict(self.__state.storage)
                 for name in list(storage["images"]):
                     del storage["images"][name]["path"]
                     del storage["images"][name]["in_storage"]
 
                 storage["downloading"] = (self.__reader.get_state() if self.__reader else None)
-
-                if self.__writer:
-                    # При загрузке файла показываем актуальную статистику вручную
-                    storage["uploading"] = self.__writer.get_state()
-                    space = self.__storage.get_space(fatal=False)
-                    if space:
-                        storage.update(dataclasses.asdict(space))
-                else:
-                    storage["uploading"] = None
+                storage["uploading"] = (self.__writer.get_state() if self.__writer else None)
 
             vd: (dict | None) = None
             if self.__state.vd:
@@ -373,7 +361,6 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
             await image.remount_rw(True)
             try:
                 await image.remove(fatal=True)
-                del self.__state.storage.images[name]
             finally:
                 await aiotools.shield_fg(image.remount_rw(False, fatal=False))
 
@@ -431,7 +418,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
                 with Inotify() as inotify:
                     for path in [
-                        *(await self.__storage.get_watchable_paths()),
+                        *self.__storage.get_watchable_paths(),
                         *self.__drive.get_watchable_paths(),
                     ]:
                         await inotify.watch(path, InotifyMask.ALL_MODIFY_EVENTS)
@@ -461,7 +448,14 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
         logger = get_logger(0)
         async with self.__state._lock:  # pylint: disable=protected-access
             try:
-                drive_state = await self.__make_init_drive_state()
+                path = self.__drive.get_image_path()
+                drive_state = _DriveState(
+                    image=((await self.__storage.make_image_by_path(path)) if path else None),
+                    cdrom=self.__drive.get_cdrom_flag(),
+                    rw=self.__drive.get_rw_flag(),
+                )
+
+                await self.__storage.reload()
 
                 if self.__state.vd is None and drive_state.image is None:
                     # Если только что включились и образ не подключен - попробовать
@@ -471,15 +465,13 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                     await self.__storage.remount_rw(False)
                     await self.__setup_initial()
 
-                storage_state = await self.__make_init_storage_state()
-
             except Exception:
                 logger.exception("Error while reloading MSD state; switching to offline")
                 self.__state.storage = None
                 self.__state.vd = None
 
             else:
-                self.__state.storage = storage_state
+                self.__state.storage = self.__storage
                 if drive_state.image:
                     # При подключенном образе виртуальный стейт заменяется реальным
                     self.__state.vd = _VirtualDriveState.from_drive_state(drive_state)
@@ -511,23 +503,3 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                     logger.exception("Can't setup initial image: ignored")
             else:
                 logger.error("Can't find initial image %r: ignored", self.__initial_image)
-
-    # =====
-
-    async def __make_init_storage_state(self) -> _StorageState:
-        images = await self.__storage.get_images()
-        space = self.__storage.get_space(fatal=True)
-        assert space
-        return _StorageState(
-            size=space.size,
-            free=space.free,
-            images=images,
-        )
-
-    async def __make_init_drive_state(self) -> _DriveState:
-        path = self.__drive.get_image_path()
-        return _DriveState(
-            image=((await self.__storage.make_image_by_path(path)) if path else None),
-            cdrom=self.__drive.get_cdrom_flag(),
-            rw=self.__drive.get_rw_flag(),
-        )
diff --git a/kvmd/plugins/msd/otg/storage.py b/kvmd/plugins/msd/otg/storage.py
index b65e3b06..e56107ff 100644
--- a/kvmd/plugins/msd/otg/storage.py
+++ b/kvmd/plugins/msd/otg/storage.py
@@ -31,8 +31,6 @@ from typing import Optional
 import aiofiles
 import aiofiles.os
 
-from ....logging import get_logger
-
 from .... import aiotools
 from .... import aiohelpers
 
@@ -59,7 +57,7 @@ class Image(_Image):
         self.__complete_path = os.path.join(self.__dir_path, f".__{file_name}.complete")
         self.__adopted = False
 
-    async def _update(self) -> None:
+    async def _reload(self) -> None:  # Only for Storage() and set_complete()
         # adopted используется в последующих проверках
         self.__adopted = await aiotools.run_async(self.__is_adopted)
         (complete, removable, (size, mod_ts)) = await asyncio.gather(
@@ -80,7 +78,7 @@ class Image(_Image):
         path = self.path
         while not os.path.ismount(path):
             path = os.path.dirname(path)
-        return (self.__storage._get_root_path() != path)
+        return (self.__storage._get_root_path() != path)  # pylint: disable=protected-access
 
     async def __is_complete(self) -> bool:
         if self.__storage:
@@ -119,6 +117,7 @@ class Image(_Image):
         assert self.__storage
         try:
             await aiofiles.os.remove(self.path)
+            self.__storage.images.pop(self.name, None)
         except FileNotFoundError:
             pass
         except Exception:
@@ -136,45 +135,60 @@ class Image(_Image):
                 await aiofiles.os.remove(self.__complete_path)
             except FileNotFoundError:
                 pass
-        await self._update()
+        await self._reload()
 
 
[email protected](frozen=True)
-class StorageSpace:
[email protected](frozen=True, eq=False)
+class _Storage:
     size: int
     free: int
+    images: dict[str, Image] = dataclasses.field(init=False)
 
 
-class Storage:
+class Storage(_Storage):
     def __init__(self, path: str, remount_cmd: list[str]) -> None:
+        super().__init__(0, 0)
         self.__path = path
         self.__remount_cmd = remount_cmd
+        self.__watchable_paths: (list[str] | None) = None
+        self.__images: (dict[str, Image] | None) = None
 
-    def _get_root_path(self) -> str:
-        return self.__path
+    @property
+    def images(self) -> dict[str, Image]:
+        assert self.__watchable_paths is not None
+        assert self.__images is not None
+        return self.__images
+
+    async def reload(self) -> None:
+        self.__watchable_paths = None
+        self.__images = {}
+
+        watchable_paths: list[str] = []
+        images: dict[str, Image] = {}
+        for (root_path, files) in (await aiotools.run_async(self.__walk)):
+            watchable_paths.append(root_path)
+            for path in files:
+                name = os.path.relpath(path, self.__path)
+                images[name] = await self.make_image_by_name(name)
+
+        await self.reload_size_only()
 
-    async def get_watchable_paths(self) -> list[str]:
-        return (await aiotools.run_async(self.__inner_get_watchable_paths))
+        self.__watchable_paths = watchable_paths
+        self.__images = images
 
-    async def get_images(self) -> dict[str, Image]:
-        return {
-            name: (await self.make_image_by_name(name))
-            for name in (await aiotools.run_async(self.__inner_get_images))
-        }
+    async def reload_size_only(self) -> None:
+        st = os.statvfs(self.__path)  # FIXME
+        object.__setattr__(self, "size", st.f_blocks * st.f_frsize)
+        object.__setattr__(self, "free", st.f_bavail * st.f_frsize)
 
-    def __inner_get_watchable_paths(self) -> list[str]:
-        return list(map(operator.itemgetter(0), self.__walk(with_files=False)))
+    def get_watchable_paths(self) -> list[str]:
+        assert self.__watchable_paths is not None
+        return list(self.__watchable_paths)
 
-    def __inner_get_images(self) -> list[str]:
-        return [
-            os.path.relpath(path, self.__path)  # == name
-            for (_, files) in self.__walk(with_files=True)
-            for path in files
-        ]
+    def __walk(self) -> list[tuple[str, list[str]]]:
+        return list(self.__inner_walk(self.__path))
 
-    def __walk(self, with_files: bool, root_path: (str | None)=None) -> Generator[tuple[str, list[str]], None, None]:
-        if root_path is None:
-            root_path = self.__path
+    def __inner_walk(self, root_path: str) -> Generator[tuple[str, list[str]], None, None]:
         files: list[str] = []
         with os.scandir(root_path) as dir_iter:
             for item in sorted(dir_iter, key=operator.attrgetter("name")):
@@ -183,8 +197,8 @@ class Storage:
                 try:
                     if item.is_dir(follow_symlinks=False):
                         item.stat()  # Проверяем, не сдохла ли смонтированная NFS
-                        yield from self.__walk(with_files, item.path)
-                    elif with_files and item.is_file(follow_symlinks=False):
+                        yield from self.__inner_walk(item.path)
+                    elif item.is_file(follow_symlinks=False):
                         files.append(item.path)
                 except Exception:
                     pass
@@ -195,7 +209,7 @@ class Storage:
     async def make_image_by_name(self, name: str) -> Image:
         assert name
         path = os.path.join(self.__path, name)
-        return (await self.__get_image(name, path, True))
+        return (await self.__make_image(name, path, True))
 
     async def make_image_by_path(self, path: str) -> Image:
         assert path
@@ -204,29 +218,19 @@ class Storage:
             name = os.path.relpath(path, self.__path)
         else:
             name = os.path.basename(path)
-        return (await self.__get_image(name, path, in_storage))
+        return (await self.__make_image(name, path, in_storage))
 
-    async def __get_image(self, name: str, path: str, in_storage: bool) -> Image:
+    async def __make_image(self, name: str, path: str, in_storage: bool) -> Image:
         assert name
         assert path
         image = Image(name, path, (self if in_storage else None))
-        await image._update()  # pylint: disable=protected-access
+        await image._reload()  # pylint: disable=protected-access
         return image
 
-    # =====
+    def _get_root_path(self) -> str:  # Only for Image()
+        return self.__path
 
-    def get_space(self, fatal: bool) -> (StorageSpace | None):
-        try:
-            st = os.statvfs(self.__path)
-        except Exception as err:
-            if fatal:
-                raise
-            get_logger().warning("Can't get free space of filesystem %s: %s", self.__path, err)
-            return None
-        return StorageSpace(
-            size=(st.f_blocks * st.f_frsize),
-            free=(st.f_bavail * st.f_frsize),
-        )
+    # =====
 
     async def remount_rw(self, rw: bool, fatal: bool=True) -> None:
         if not (await aiohelpers.remount("MSD", self.__remount_cmd, rw)):