refactoring

2 files changed, 109 insertions, 124 deletions

diff --git a/kvmd/plugins/msd/otg/__init__.py b/kvmd/plugins/msd/otg/__init__.py
index 4a8376d0..e8805313 100644
--- a/kvmd/plugins/msd/otg/__init__.py
+++ b/kvmd/plugins/msd/otg/__init__.py
@@ -20,7 +20,6 @@
 # ========================================================================== #
 
 
-import os
 import asyncio
 import contextlib
 import dataclasses
@@ -57,37 +56,15 @@ from .. import BaseMsd
 from .. import MsdFileReader
 from .. import MsdFileWriter
 
+from .storage import Image
 from .storage import Storage
 from .drive import Drive
 
 
 # =====
 @dataclasses.dataclass(frozen=True)
-class _DriveImage:
-    name: str
-    path: str
-    complete: bool
-    in_storage: bool
-    size: int = dataclasses.field(default=0)
-    mod_ts: float = dataclasses.field(default=0)
-
-    @property
-    def exists(self) -> bool:  # Not exposed as a field
-        return os.path.exists(self.path)
-
-    def __post_init__(self) -> None:
-        try:
-            st = os.stat(self.path)
-        except Exception as err:
-            get_logger().warning("Can't stat() file %s: %s", self.path, err)
-        else:
-            object.__setattr__(self, "size", st.st_size)
-            object.__setattr__(self, "mod_ts", st.st_mtime)
-
-
[email protected](frozen=True)
 class _DriveState:
-    image: (_DriveImage | None)
+    image: (Image | None)
     cdrom: bool
     rw: bool
 
@@ -96,13 +73,13 @@ class _DriveState:
 class _StorageState:
     size: int
     free: int
-    images: dict[str, _DriveImage]
+    images: dict[str, Image]
 
 
 # =====
 @dataclasses.dataclass
 class _VirtualDriveState:
-    image: (_DriveImage | None)
+    image: (Image | None)
     connected: bool
     cdrom: bool
     rw: bool
@@ -124,8 +101,8 @@ class _State:
         self.storage: (_StorageState | None) = None
         self.vd: (_VirtualDriveState | None) = None
 
-        self._lock = asyncio.Lock()
         self._region = aiotools.AioExclusiveRegion(MsdIsBusyError)
+        self._lock = asyncio.Lock()
 
     @contextlib.asynccontextmanager
     async def busy(self, check_online: bool=True) -> AsyncGenerator[None, None]:
@@ -276,16 +253,12 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
     ) -> None:
 
         async with self.__state.busy():
-            assert self.__state.storage
             assert self.__state.vd
-
-            self.__check_disconnected()
+            self.__state_check_disconnected()
 
             if name is not None:
                 if name:
-                    image = self.__get_image(name)
-                    assert image.in_storage
-                    self.__state.vd.image = image
+                    self.__state.vd.image = self.__state_get_storage_image(name)
                 else:
                     self.__state.vd.image = None
 
@@ -304,16 +277,16 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
         async with self.__state.busy():
             assert self.__state.vd
             if connected:
-                self.__check_disconnected()
+                self.__state_check_disconnected()
 
                 if self.__state.vd.image is None:
                     raise MsdImageNotSelected()
 
-                assert self.__state.vd.image.in_storage
-
-                if not self.__state.vd.image.exists:
+                if not self.__state.vd.image.exists():
                     raise MsdUnknownImageError()
 
+                assert self.__state.vd.image.in_storage
+
                 self.__drive.set_rw_flag(self.__state.vd.rw)
                 self.__drive.set_cdrom_flag(self.__state.vd.cdrom)
                 if self.__state.vd.rw:
@@ -321,7 +294,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                 self.__drive.set_image_path(self.__state.vd.image.path)
 
             else:
-                self.__check_connected()
+                self.__state_check_connected()
                 self.__drive.set_image_path("")
                 await self.__remount_rw(False, fatal=False)
 
@@ -334,21 +307,14 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                 try:
                     async with self.__state._lock:  # pylint: disable=protected-access
                         self.__notifier.notify()
-                        assert self.__state.storage
-                        assert self.__state.vd
-
-                        self.__check_disconnected()
-
-                        image = self.__get_image(name)
-
+                        self.__state_check_disconnected()
+                        image = self.__state_get_storage_image(name)
                         self.__reader = await MsdFileReader(
                             notifier=self.__notifier,
                             path=image.path,
                             chunk_size=self.__read_chunk_size,
                         ).open()
-
                     yield self.__reader
-
                 finally:
                     await aiotools.shield_fg(self.__close_reader())
         finally:
@@ -358,25 +324,23 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
     async def write_image(self, name: str, size: int, remove_incomplete: (bool | None)) -> AsyncGenerator[MsdFileWriter, None]:
         try:
             async with self.__state._region:  # pylint: disable=protected-access
-                path: str = ""
+                image: (Image | None) = None
                 try:
                     async with self.__state._lock:  # pylint: disable=protected-access
                         self.__notifier.notify()
                         assert self.__state.storage
-                        assert self.__state.vd
-
-                        self.__check_disconnected()
+                        self.__state_check_disconnected()
 
-                        path = self.__storage.get_image_path(name)
-                        if name in self.__state.storage.images or os.path.exists(path):
+                        image = self.__storage.get_image_by_name(name)
+                        if image.name in self.__state.storage.images or image.exists():
                             raise MsdImageExistsError()
 
                         await self.__remount_rw(True)
-                        self.__storage.set_image_complete(name, False)
+                        self.__storage.set_image_complete(image, False)
 
                         self.__writer = await MsdFileWriter(
                             notifier=self.__notifier,
-                            path=path,
+                            path=image.path,
                             file_size=size,
                             sync_size=self.__sync_chunk_size,
                             chunk_size=self.__write_chunk_size,
@@ -384,15 +348,11 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
                     self.__notifier.notify()
                     yield self.__writer
-                    self.__storage.set_image_complete(name, self.__writer.is_complete())
+                    self.__storage.set_image_complete(image, self.__writer.is_complete())
 
                 finally:
-                    if remove_incomplete and self.__writer and not self.__writer.is_complete():
-                        # Можно сперва удалить файл, потом закрыть его
-                        try:
-                            os.remove(path)
-                        except Exception:
-                            pass
+                    if image and remove_incomplete and self.__writer and not self.__writer.is_complete():
+                        self.__storage.remove_image(image, fatal=False)
                     try:
                         await aiotools.shield_fg(self.__close_writer())
                     finally:
@@ -407,38 +367,38 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
         async with self.__state.busy():
             assert self.__state.storage
             assert self.__state.vd
+            self.__state_check_disconnected()
 
-            self.__check_disconnected()
-
-            image = self.__get_image(name)
-            assert image.in_storage
+            image = self.__state_get_storage_image(name)
 
             if self.__state.vd.image == image:
                 self.__state.vd.image = None
             del self.__state.storage.images[name]
 
             await self.__remount_rw(True)
-            os.remove(image.path)
-            self.__storage.set_image_complete(name, False)
-            await self.__remount_rw(False, fatal=False)
+            try:
+                self.__storage.remove_image(image, fatal=True)
+            finally:
+                await self.__remount_rw(False, fatal=False)
 
     # =====
 
-    def __check_connected(self) -> None:
+    def __state_check_connected(self) -> None:
         assert self.__state.vd
         if not (self.__state.vd.connected or self.__drive.get_image_path()):
             raise MsdDisconnectedError()
 
-    def __check_disconnected(self) -> None:
+    def __state_check_disconnected(self) -> None:
         assert self.__state.vd
         if self.__state.vd.connected or self.__drive.get_image_path():
             raise MsdConnectedError()
 
-    def __get_image(self, name: str) -> _DriveImage:
+    def __state_get_storage_image(self, name: str) -> Image:
         assert self.__state.storage
         image = self.__state.storage.images.get(name)
-        if image is None or not image.exists:
+        if image is None or not image.exists():
             raise MsdUnknownImageError()
+        assert image.in_storage
         return image
 
     async def __close_reader(self) -> None:
@@ -507,10 +467,12 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                     await self.__setup_initial()
 
                 storage_state = self.__get_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
                 if drive_state.image:
@@ -521,10 +483,8 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                         # Если раньше MSD был отключен
                         self.__state.vd = _VirtualDriveState.from_drive_state(drive_state)
 
-                    if (
-                        self.__state.vd.image
-                        and (not self.__state.vd.image.in_storage or not self.__state.vd.image.exists)
-                    ):
+                    image = self.__state.vd.image
+                    if image and (not image.in_storage or not image.exists()):
                         # Если только что отключили ручной образ вне хранилища или ранее выбранный образ был удален
                         self.__state.vd.image = None
 
@@ -535,13 +495,13 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
     async def __setup_initial(self) -> None:
         if self.__initial_image:
             logger = get_logger(0)
-            path = self.__storage.get_image_path(self.__initial_image)
-            if os.path.exists(path):
+            image = self.__storage.get_image_by_name(self.__initial_image)
+            if image.exists():
                 logger.info("Setting up initial image %r ...", self.__initial_image)
                 try:
                     self.__drive.set_rw_flag(False)
                     self.__drive.set_cdrom_flag(self.__initial_cdrom)
-                    self.__drive.set_image_path(path)
+                    self.__drive.set_image_path(image.path)
                 except Exception:
                     logger.exception("Can't setup initial image: ignored")
             else:
@@ -550,14 +510,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
     # =====
 
     def __get_storage_state(self) -> _StorageState:
-        images: dict[str, _DriveImage] = {}
-        for name in self.__storage.get_images():
-            images[name] = _DriveImage(
-                name=name,
-                path=self.__storage.get_image_path(name),
-                complete=self.__storage.is_image_complete(name),
-                in_storage=True,
-            )
+        images = self.__storage.get_images()
         space = self.__storage.get_space(fatal=True)
         assert space
         return _StorageState(
@@ -567,19 +520,9 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
         )
 
     def __get_drive_state(self) -> _DriveState:
-        image: (_DriveImage | None) = None
         path = self.__drive.get_image_path()
-        if path:
-            name = os.path.basename(path)
-            in_storage = self.__storage.is_image_path_in_storage(path)
-            image = _DriveImage(
-                name=name,
-                path=path,
-                complete=(self.__storage.is_image_complete(name) if in_storage else True),
-                in_storage=in_storage,
-            )
         return _DriveState(
-            image=image,
+            image=(self.__storage.get_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 28ab0fa4..beb7f620 100644
--- a/kvmd/plugins/msd/otg/storage.py
+++ b/kvmd/plugins/msd/otg/storage.py
@@ -28,6 +28,30 @@ from ....logging import get_logger
 
 # =====
 @dataclasses.dataclass(frozen=True)
+class Image:
+    name: str
+    path: str
+
+    complete: bool = dataclasses.field(compare=False)
+    in_storage: bool = dataclasses.field(compare=False)
+
+    size: int = dataclasses.field(default=0, compare=False)
+    mod_ts: float = dataclasses.field(default=0, compare=False)
+
+    def exists(self) -> bool:
+        return os.path.exists(self.path)
+
+    def __post_init__(self) -> None:
+        try:
+            st = os.stat(self.path)
+        except Exception:
+            pass
+        else:
+            object.__setattr__(self, "size", st.st_size)
+            object.__setattr__(self, "mod_ts", st.st_mtime)
+
+
[email protected](frozen=True)
 class StorageSpace:
     size: int
     free: int
@@ -42,34 +66,52 @@ class Storage:
     def get_watchable_paths(self) -> list[str]:
         return [self.__images_path, self.__meta_path]
 
-    def get_images(self) -> list[str]:
-        images: list[str] = []
-        for name in os.listdir(self.__images_path):
-            path = os.path.join(self.__images_path, name)
-            if os.path.exists(path):
-                try:
-                    if os.path.getsize(path) >= 0:
-                        images.append(name)
-                except Exception:
-                    pass
-        return images
-
-    def get_image_path(self, name: str) -> str:
-        return os.path.join(self.__images_path, name)
-
-    def is_image_path_in_storage(self, path: str) -> bool:
-        return (os.path.dirname(path) == self.__images_path)
-
-    def is_image_complete(self, name: str) -> bool:
-        return os.path.exists(os.path.join(self.__meta_path, name + ".complete"))
-
-    def set_image_complete(self, name: str, flag: bool) -> None:
-        path = os.path.join(self.__meta_path, name + ".complete")
+    def get_images(self) -> dict[str, Image]:
+        return {
+            name: self.get_image_by_name(name)
+            for name in os.listdir(self.__images_path)
+        }
+
+    def get_image_by_name(self, name: str) -> Image:
+        assert name
+        path = os.path.join(self.__images_path, name)
+        return self.__get_image(name, path)
+
+    def get_image_by_path(self, path: str) -> Image:
+        assert path
+        name = os.path.basename(path)
+        return self.__get_image(name, path)
+
+    def __get_image(self, name: str, path: str) -> Image:
+        assert name
+        assert path
+        complete = True
+        in_storage = (os.path.dirname(path) == self.__images_path)
+        if in_storage:
+            complete = os.path.exists(os.path.join(self.__meta_path, name + ".complete"))
+        return Image(name, path, complete, in_storage)
+
+    def remove_image(self, image: Image, fatal: bool) -> None:
+        assert image.in_storage
+        try:
+            os.remove(image.path)
+        except FileNotFoundError:
+            pass
+        except Exception:
+            if fatal:
+                raise
+        self.set_image_complete(image, False)
+
+    def set_image_complete(self, image: Image, flag: bool) -> None:
+        assert image.in_storage
+        path = os.path.join(self.__meta_path, image.name + ".complete")
         if flag:
             open(path, "w").close()  # pylint: disable=consider-using-with
         else:
-            if os.path.exists(path):
+            try:
                 os.remove(path)
+            except FileNotFoundError:
+                pass
 
     def get_space(self, fatal: bool) -> (StorageSpace | None):
         try: