added FileLikeFileStreamReader

Author: Andrei Neagu

packages/service-library/src/servicelib/zip_stream/__init__.py

@@ -1,4 +1,5 @@
 from ._constants import DEFAULT_READ_CHUNK_SIZE, MIN_MULTIPART_UPLOAD_CHUNK_SIZE
+from ._file_like import FileLikeFileStreamReader
 from ._input import DiskStreamReader
 from ._output import DiskStreamWriter
 from ._types import (
@@ -16,6 +17,7 @@
     "DEFAULT_READ_CHUNK_SIZE",
     "DiskStreamReader",
     "DiskStreamWriter",
+    "FileLikeFileStreamReader",
     "FileSize",
     "FileStream",
     "FileStreamCallable",

packages/service-library/src/servicelib/zip_stream/_file_like.py

@@ -0,0 +1,23 @@
+from ._types import FileStream
+
+
+class FileLikeFileStreamReader:
+    def __init__(self, file_stream: FileStream):
+        self.file_stream = file_stream
+        self._buffer = bytearray()
+        self._async_iterator = self._get_iterator()
+
+    async def _get_iterator(self):
+        async for chunk in self.file_stream:
+            yield chunk
+
+    async def read(self, size: int) -> bytes:
+        while len(self._buffer) < size:
+            try:
+                chunk = await anext(self._async_iterator)
+                self._buffer.extend(chunk)
+            except StopAsyncIteration:
+                break  # End of file
+
+        result, self._buffer = self._buffer[:size], self._buffer[size:]
+        return bytes(result)

packages/service-library/src/servicelib/zip_stream/_output.py

@@ -1,6 +1,7 @@
 from pathlib import Path
 
 import aiofiles
+from servicelib.zip_stream._file_like import FileLikeFileStreamReader
 
 from ._types import FileStream
 
@@ -9,8 +10,20 @@ class DiskStreamWriter:
     def __init__(self, destination_path: Path):
         self.destination_path = destination_path
 
-    async def write_stream(self, stream: FileStream) -> None:
+    async def write_from_stream(self, stream: FileStream) -> None:
         async with aiofiles.open(self.destination_path, "wb") as f:
             async for chunk in stream:
                 await f.write(chunk)
                 await f.flush()
+
+    async def write_from_file_like(
+        self, file_like_reader: FileLikeFileStreamReader
+    ) -> None:
+        async with aiofiles.open(self.destination_path, "wb") as f:
+            while True:
+                chunk = await file_like_reader.read(100)
+                if not chunk:
+                    break
+
+                await f.write(chunk)
+                await f.flush()

packages/service-library/tests/test_zip_stream.py

@@ -21,6 +21,7 @@
     ArchiveEntries,
     DiskStreamReader,
     DiskStreamWriter,
+    FileLikeFileStreamReader,
     get_zip_archive_stream,
 )
 
@@ -92,12 +93,14 @@ def _progress_cb(*args, **kwargs) -> None:
     return mocker.Mock(side_effect=_progress_cb)
 
 
+@pytest.mark.parametrize("use_file_like", [True, False])
 async def test_get_zip_archive_stream(
     mocked_progress_bar_cb: Mock,
     prepare_content: None,
     local_files_dir: Path,
     local_archive_path: Path,
     local_unpacked_archive: Path,
+    use_file_like: bool,
 ):
     # 1. generate archive form soruces
     archive_files: ArchiveEntries = []
@@ -113,10 +116,15 @@ async def test_get_zip_archive_stream(
         progress_report_cb=mocked_progress_bar_cb,
         description="root_bar",
     ) as root:
-        await writer.write_stream(
-            get_zip_archive_stream(archive_files, progress_bar=root, chunk_size=1024)
+        file_stream = get_zip_archive_stream(
+            archive_files, progress_bar=root, chunk_size=1024
         )
 
+        if use_file_like:
+            await writer.write_from_file_like(FileLikeFileStreamReader(file_stream))
+        else:
+            await writer.write_from_stream(file_stream)
+
     # 2. extract archive using exiting tools
     await unarchive_dir(local_archive_path, local_unpacked_archive)