hojichar.utils.async_handlers

  1from __future__ import annotations
  2
  3import asyncio
  4import itertools
  5import logging
  6import queue
  7import threading
  8from concurrent.futures import ThreadPoolExecutor
  9from pathlib import Path
 10from typing import (
 11    Any,
 12    AsyncGenerator,
 13    AsyncIterable,
 14    Awaitable,
 15    Callable,
 16    Generic,
 17    Iterable,
 18    Iterator,
 19    TextIO,
 20    TypeVar,
 21    cast,
 22)
 23
 24T = TypeVar("T")
 25logger = logging.getLogger(__name__)
 26
 27
 28class _AsyncIteratorError:
 29    def __init__(self, error: BaseException):
 30        self.error = error
 31
 32
 33_ASYNC_ITERATOR_END = object()
 34
 35
 36class AsyncToSyncIterator(Iterator[T], Generic[T]):
 37    """Consume an async iterable from synchronous code.
 38
 39    A single background thread owns the event loop for the lifetime of the iterator. Use
 40    this class as a context manager when iteration may stop before the source is exhausted.
 41    """
 42
 43    def __init__(
 44        self,
 45        source_stream: AsyncIterable[T],
 46        *,
 47        buffer_size: int = 128,
 48        finalizer: Callable[[], Awaitable[None]] | None = None,
 49    ) -> None:
 50        if buffer_size < 1:
 51            raise ValueError("buffer_size must be at least 1")
 52
 53        self._source_stream = source_stream
 54        self._queue: queue.Queue[Any] = queue.Queue(maxsize=buffer_size)
 55        self._finalizer = finalizer
 56        self._stop_event = threading.Event()
 57        self._started_event = threading.Event()
 58        self._thread: threading.Thread | None = None
 59        self._loop: asyncio.AbstractEventLoop | None = None
 60        self._consumer_task: asyncio.Task[None] | None = None
 61        self._closed = False
 62
 63    def __iter__(self) -> "AsyncToSyncIterator[T]":
 64        return self
 65
 66    def __next__(self) -> T:
 67        if self._closed:
 68            raise StopIteration
 69        self._start()
 70
 71        item = self._queue.get()
 72        if item is _ASYNC_ITERATOR_END:
 73            self.close()
 74            raise StopIteration
 75        if isinstance(item, _AsyncIteratorError):
 76            self.close()
 77            raise item.error
 78        return cast(T, item)
 79
 80    def __enter__(self) -> "AsyncToSyncIterator[T]":
 81        return self
 82
 83    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
 84        self.close()
 85
 86    def close(self) -> None:
 87        if self._closed:
 88            return
 89        self._closed = True
 90        self._stop_event.set()
 91        thread_was_started = self._thread is not None
 92
 93        if not thread_was_started:
 94            # Closing an unconsumed iterator must still close its source and run its finalizer.
 95            self._thread = threading.Thread(
 96                target=self._thread_main,
 97                name="hojichar-async-to-sync",
 98                daemon=True,
 99            )
100            self._thread.start()
101            self._started_event.wait()
102
103        while True:
104            try:
105                self._queue.get_nowait()
106            except queue.Empty:
107                break
108
109        if (
110            thread_was_started
111            and self._thread is not None
112            and self._thread.is_alive()
113            and self._loop is not None
114            and not self._loop.is_closed()
115            and self._consumer_task is not None
116        ):
117            try:
118                self._loop.call_soon_threadsafe(self._consumer_task.cancel)
119            except RuntimeError:
120                # The background loop can close between is_closed() and scheduling.
121                pass
122
123        if self._thread is not None:
124            self._thread.join(timeout=5.0)
125            if self._thread.is_alive():
126                logger.warning("Async-to-sync iterator did not stop within 5 seconds")
127
128    def _start(self) -> None:
129        if self._thread is not None:
130            return
131        try:
132            asyncio.get_running_loop()
133        except RuntimeError:
134            pass
135        else:
136            raise RuntimeError(
137                "Cannot synchronously consume an async iterable from a running event loop"
138            )
139
140        self._thread = threading.Thread(
141            target=self._thread_main,
142            name="hojichar-async-to-sync",
143            daemon=True,
144        )
145        self._thread.start()
146        self._started_event.wait()
147
148    def _thread_main(self) -> None:
149        try:
150            asyncio.run(self._consume())
151        except BaseException as error:
152            if not self._stop_event.is_set():
153                self._put(_AsyncIteratorError(error))
154            elif not isinstance(error, asyncio.CancelledError):
155                logger.error("Failed to close async-to-sync iterator", exc_info=True)
156        finally:
157            self._started_event.set()
158            if not self._stop_event.is_set():
159                self._put(_ASYNC_ITERATOR_END)
160
161    async def _consume(self) -> None:
162        self._loop = asyncio.get_running_loop()
163        self._consumer_task = asyncio.current_task()
164        self._started_event.set()
165        iterator = None
166
167        try:
168            iterator = self._source_stream.__aiter__()
169            while not self._stop_event.is_set():
170                try:
171                    item = await iterator.__anext__()
172                except StopAsyncIteration:
173                    break
174                if not self._put(item):
175                    break
176        finally:
177            try:
178                aclose = getattr(iterator, "aclose", None)
179                if aclose is not None:
180                    await aclose()
181            finally:
182                if self._finalizer is not None:
183                    await self._finalizer()
184
185    def _put(self, item: Any) -> bool:
186        while not self._stop_event.is_set():
187            try:
188                self._queue.put(item, timeout=0.05)
189                return True
190            except queue.Full:
191                continue
192        return False
193
194
195def handle_stream_as_async(
196    source_stream: Iterable[T] | AsyncIterable[T],
197    chunk_size: int = 1000,
198    executor: ThreadPoolExecutor | None = None,
199) -> AsyncGenerator[T, None]:
200    """
201    Convert a synchronous iterable to an asynchronous generator
202    with a specified chunk size.
203
204    Args:
205        source_stream (Iterable[T]): The synchronous iterable to convert.
206        chunk_size (int): The number of items to yield at a time.
207    """
208    if isinstance(source_stream, AsyncIterable):
209        return source_stream  # type: ignore[return-value]
210    stream = iter(source_stream)
211
212    async def sync_to_async() -> AsyncGenerator[T, None]:
213        loop = asyncio.get_running_loop()
214        while True:
215            chunk = await loop.run_in_executor(
216                executor, lambda: list(itertools.islice(stream, chunk_size))
217            )
218            if not chunk:
219                break
220            for item in chunk:
221                yield item
222
223    return sync_to_async()
224
225
226def handle_async_stream_as_sync(
227    source_stream: AsyncIterable[T],
228    *,
229    buffer_size: int = 128,
230    finalizer: Callable[[], Awaitable[None]] | None = None,
231) -> AsyncToSyncIterator[T]:
232    """Convert an async iterable to a synchronous, closeable iterator.
233
234    The iterator owns a background event-loop thread. Fully consuming it closes the source
235    automatically. Use it as a context manager when the consumer may stop early.
236    """
237    return AsyncToSyncIterator(
238        source_stream,
239        buffer_size=buffer_size,
240        finalizer=finalizer,
241    )
242
243
244async def write_stream_to_file(
245    stream: AsyncGenerator[str, None],
246    output_path: Path | str,
247    *,
248    chunk_size: int = 1000,
249    delimiter: str = "\n",
250) -> None:
251    """
252    Write an asynchronous stream of strings to a file.
253    To lessen overhead with file I/O, it writes in chunks.
254    """
255    loop = asyncio.get_running_loop()
256    with open(output_path, "w", encoding="utf-8") as f:
257        chunk = []
258        async for line in stream:
259            chunk.append(line)
260            if len(chunk) >= chunk_size:
261                await loop.run_in_executor(None, f.writelines, [s + delimiter for s in chunk])
262                chunk = []
263        if chunk:
264            await loop.run_in_executor(None, f.writelines, [s + delimiter for s in chunk])
265            chunk = []
266        await loop.run_in_executor(None, f.flush)
267
268
269async def fileout_from_async_iter(
270    fp: TextIO, iter: AsyncIterable[str], buffer_size: int = 128
271) -> None:
272    buffer = []
273    async for line in iter:
274        buffer.append(line + "\n")
275        if len(buffer) >= buffer_size:
276            await asyncio.to_thread(fp.write, "".join(buffer))
277            buffer.clear()
278    await asyncio.to_thread(fp.writelines, buffer)
279    buffer.clear()