caffeinism-utils 0.0.169__tar.gz → 0.0.171__tar.gz

{caffeinism_utils-0.0.169 → caffeinism_utils-0.0.171}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: caffeinism-utils
-Version: 0.0.169
+Version: 0.0.171
 Summary: 
 Author: Kim Minjong
 Author-email: make.dirty.code@gmail.com

{caffeinism_utils-0.0.169 → caffeinism_utils-0.0.171}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "caffeinism-utils"
-version = "0.0.169"
+version = "0.0.171"
 description = ""
 authors = ["Kim Minjong <make.dirty.code@gmail.com>"]
 readme = "README.md"

{caffeinism_utils-0.0.169 → caffeinism_utils-0.0.171}/src/caffeinism_utils/aiteration.py

@@ -1,6 +1,6 @@
 import asyncio
 import time
-from typing import AsyncGenerator, Generic, TypeVar
+from typing import AsyncGenerator, AsyncIterable, Generic, TypeVar
 
 from .asyncio import run_in_threadpool
 from .utils import DummyStopIteration, next_without_stop_iteration
@@ -26,15 +26,35 @@ async def rate_limit_iterator(aiterator, iters_per_second):
 
 T = TypeVar("T")
 
+__CLOSE = object()
+
 
 class StreamQueue(Generic[T]):
     def __init__(self):
         self.queue = asyncio.Queue[T]()
 
-    def put(self, data: T) -> T:
+    def put(self, data: T) -> None:
+        return self.queue.put(data)
+
+    async def __aiter__(self) -> AsyncGenerator[T]:
+        while (it := await self.queue.get()) is not __CLOSE:
+            yield it
+
+    async def close(self):
+        await self.queue.put(__CLOSE)
+
+
+class StreamIteratorQueue(Generic[T]):
+    def __init__(self):
+        self.queue = asyncio.Queue[AsyncIterable[T]]()
+
+    def put(self, data: AsyncIterable[T]) -> None:
         return self.queue.put(data)
 
     async def __aiter__(self) -> AsyncGenerator[T]:
-        while (iterator := await self.queue.get()) is not None:
+        while (iterator := await self.queue.get()) is not __CLOSE:
             async for it in iterator:
                 yield it
+
+    async def close(self):
+        await self.queue.put(__CLOSE)

{caffeinism_utils-0.0.169 → caffeinism_utils-0.0.171}/src/caffeinism_utils/prefetch.py

@@ -103,7 +103,7 @@ async def _aprefetch_aiterator(
     while True:
         if prefetched is not None:
             try:
-                ret = await prefetched
+                ret = await asyncio.shield(prefetched)
             except StopAsyncIteration:
                 break