refactored msd reader api

kvmd/apps/kvmd/api/msd.py

@@ -87,12 +87,10 @@ class MsdApi:
     @exposed_http("GET", "/msd/read")
     async def __read_handler(self, request: Request) -> StreamResponse:
         name = valid_msd_image_name(request.query.get("image"))
-        async with self.__msd.read_image(name) as size:
+        async with self.__msd.read_image(name) as reader:
+            size = reader.get_total_size()
             response = await start_streaming(request, "application/octet-stream", size, name)
-            while True:
+            async for chunk in reader.read_chunked():
-                chunk = await self.__msd.read_image_chunk()
-                if not chunk:
-                    return response
                 await response.write(chunk)
             return response
 

kvmd/plugins/msd/__init__.py

@@ -104,6 +104,19 @@ class MsdRwNotSupported(MsdOperationError):
 
 
 # =====
+class BaseMsdReader:
+    def get_state(self) -> Dict:
+        raise NotImplementedError()
+
+    def get_total_size(self) -> int:
+        raise NotImplementedError()
+
+    async def read_chunked(self) -> AsyncGenerator[bytes, None]:
+        if self is not None:  # XXX: Vulture and pylint hack
+            raise NotImplementedError()
+        yield
+
+
 class BaseMsd(BasePlugin):
     async def get_state(self) -> Dict:
         raise NotImplementedError()
@@ -134,14 +147,11 @@ class BaseMsd(BasePlugin):
         raise NotImplementedError()
 
     @contextlib.asynccontextmanager
-    async def read_image(self, name: str) -> AsyncGenerator[int, None]:
+    async def read_image(self, name: str) -> AsyncGenerator[BaseMsdReader, None]:
         _ = name
         if self is not None:  # XXX: Vulture and pylint hack
             raise NotImplementedError()
-        yield 1
+        yield BaseMsdReader()
-
-    async def read_image_chunk(self) -> bytes:
-        raise NotImplementedError()
 
     @contextlib.asynccontextmanager
     async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:
@@ -159,7 +169,7 @@ class BaseMsd(BasePlugin):
         raise NotImplementedError()
 
 
-class MsdImageReader:  # pylint: disable=too-many-instance-attributes
+class MsdFileReader(BaseMsdReader):  # pylint: disable=too-many-instance-attributes
     def __init__(self, notifier: aiotools.AioNotifier, path: str, chunk_size: int) -> None:
         self.__notifier = notifier
         self.__name = os.path.basename(path)
@@ -171,10 +181,6 @@ class MsdImageReader:  # pylint: disable=too-many-instance-attributes
         self.__readed = 0
         self.__tick = 0.0
 
-    def get_size(self) -> int:
-        assert self.__file is not None
-        return self.__file_size
-
     def get_state(self) -> Dict:
         return {
             "name": self.__name,
@@ -182,16 +188,17 @@ class MsdImageReader:  # pylint: disable=too-many-instance-attributes
             "readed": self.__readed,
         }
 
-    async def open(self) -> "MsdImageReader":
+    def get_total_size(self) -> int:
-        assert self.__file is None
-        get_logger(1).info("Reading %r image from MSD ...", self.__name)
-        self.__file_size = os.stat(self.__path).st_size
-        self.__file = await aiofiles.open(self.__path, mode="rb")  # type: ignore
-        return self
-
-    async def read(self) -> bytes:
         assert self.__file is not None
+        return self.__file_size
+
+    async def read_chunked(self) -> AsyncGenerator[bytes, None]:
+        assert self.__file is not None
+        while True:
             chunk = await self.__file.read(self.__chunk_size)  # type: ignore
+            if not chunk:
+                break
+
             self.__readed += len(chunk)
 
             now = time.monotonic()
@@ -199,7 +206,14 @@ class MsdImageReader:  # pylint: disable=too-many-instance-attributes
                 self.__tick = now
                 await self.__notifier.notify()
 
-        return chunk
+            yield chunk
+
+    async def open(self) -> "MsdFileReader":
+        assert self.__file is None
+        get_logger(1).info("Reading %r image from MSD ...", self.__name)
+        self.__file_size = os.stat(self.__path).st_size
+        self.__file = await aiofiles.open(self.__path, mode="rb")  # type: ignore
+        return self
 
     async def close(self) -> None:
         assert self.__file is not None

kvmd/plugins/msd/disabled.py

@@ -29,6 +29,7 @@ from typing import Optional
 from ... import aiotools
 
 from . import MsdOperationError
+from . import BaseMsdReader
 from . import BaseMsd
 
 
@@ -77,13 +78,10 @@ class Plugin(BaseMsd):
         raise MsdDisabledError()
 
     @contextlib.asynccontextmanager
-    async def read_image(self, name: str) -> AsyncGenerator[int, None]:
+    async def read_image(self, name: str) -> AsyncGenerator[BaseMsdReader, None]:
         if self is not None:  # XXX: Vulture and pylint hack
             raise MsdDisabledError()
-        yield 1
+        yield BaseMsdReader()
-
-    async def read_image_chunk(self) -> bytes:
-        raise MsdDisabledError()
 
     @contextlib.asynccontextmanager
     async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:

kvmd/plugins/msd/otg/__init__.py

@@ -57,7 +57,7 @@ from .. import MsdImageNotSelected
 from .. import MsdUnknownImageError
 from .. import MsdImageExistsError
 from .. import BaseMsd
-from .. import MsdImageReader
+from .. import MsdFileReader
 from .. import MsdImageWriter
 
 from . import fs
@@ -165,7 +165,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
         self.__drive = Drive(gadget, instance=0, lun=0)
 
-        self.__reader: Optional[MsdImageReader] = None
+        self.__reader: Optional[MsdFileReader] = None
         self.__writer: Optional[MsdImageWriter] = None
 
         self.__notifier = aiotools.AioNotifier()
@@ -329,7 +329,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
             self.__state.vd.connected = connected
 
     @contextlib.asynccontextmanager
-    async def read_image(self, name: str) -> AsyncGenerator[int, None]:
+    async def read_image(self, name: str) -> AsyncGenerator[MsdFileReader, None]:
         try:
             async with self.__state._region:  # pylint: disable=protected-access
                 try:
@@ -345,23 +345,19 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                         if name not in self.__state.storage.images or not os.path.exists(path):
                             raise MsdUnknownImageError()
 
-                        self.__reader = await MsdImageReader(
+                        self.__reader = await MsdFileReader(
                             notifier=self.__notifier,
                             path=path,
                             chunk_size=self.__read_chunk_size,
                         ).open()
 
-                    yield self.__reader.get_size()
+                    yield self.__reader
 
                 finally:
                     await self.__close_reader()
         finally:
             await self.__notifier.notify()
 
-    async def read_image_chunk(self) -> bytes:
-        assert self.__reader
-        return (await self.__reader.read())
-
     @contextlib.asynccontextmanager
     async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:
         try:

kvmd/plugins/msd/relay/__init__.py

@@ -50,6 +50,7 @@ from .. import MsdDisconnectedError
 from .. import MsdMultiNotSupported
 from .. import MsdCdromNotSupported
 from .. import MsdRwNotSupported
+from .. import BaseMsdReader
 from .. import BaseMsd
 from .. import MsdImageWriter
 
@@ -218,15 +219,11 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                 self.__connected = connected
 
     @contextlib.asynccontextmanager
-    async def read_image(self, name: str) -> AsyncGenerator[int, None]:
+    async def read_image(self, name: str) -> AsyncGenerator[BaseMsdReader, None]:
         async with self.__working():
             if self is not None:  # XXX: Vulture and pylint hack
                 raise MsdMultiNotSupported()
-        yield 1
+        yield BaseMsdReader()
-
-    async def read_image_chunk(self) -> bytes:
-        async with self.__working():
-            raise MsdMultiNotSupported()
 
     @contextlib.asynccontextmanager
     async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]: