diff --git a/src/athreading/iterator.py b/src/athreading/iterator.py index f0ff5a4..57d871c 100644 --- a/src/athreading/iterator.py +++ b/src/athreading/iterator.py @@ -10,7 +10,7 @@ import threading from collections.abc import Callable, Iterator from concurrent.futures import ThreadPoolExecutor -from typing import TYPE_CHECKING, Generic, Optional, TypeVar, Union, cast +from typing import TYPE_CHECKING, Generic, Optional, TypeVar, Union from athreading.aliases import AsyncIteratorContext @@ -22,12 +22,21 @@ if TYPE_CHECKING: from types import TracebackType - _ParamsT = ParamSpec("_ParamsT") _YieldT = TypeVar("_YieldT") _T = TypeVar("_T") +class _Ok(Generic[_T]): + def __init__(self, value: _T): + self.value = value + + +class _Err(Generic[_T]): + def __init__(self, error: _T): + self.error = error + + @overload def iterate( fn: None = None, @@ -144,7 +153,9 @@ def __init__( """ self._yield_semaphore = asyncio.Semaphore(0) self._done_event = threading.Event() - self._queue: queue.Queue[_Result[_YieldT]] = queue.Queue(buffer_maxsize or 0) + self._queue: queue.Queue[_Ok[_YieldT] | _Err[Exception]] = queue.Queue( + buffer_maxsize or 0 + ) self._iterator = iterator self._executor = executor self._worker_future: Optional[asyncio.Future[None]] = None @@ -178,17 +189,16 @@ async def __anext__(self) -> _YieldT: await self._yield_semaphore.acquire() if not self._queue.empty(): result = self._queue.get(False) - if result.error is not None: + if isinstance(result, _Err): raise result.error - else: - return cast(_YieldT, result.value) + return result.value raise StopAsyncIteration def __worker_threadsafe(self) -> None: """Stream the synchronous iterator to the queue and notify the async thread.""" try: for item in self._iterator: - self._queue.put(_Result(value=item, error=None)) + self._queue.put(_Ok(item)) self._loop.call_soon_threadsafe(self._yield_semaphore.release) while self._queue.full() and not self._done_event.is_set(): @@ -197,8 +207,8 @@ def __worker_threadsafe(self) -> None: if self._done_event.is_set(): break - except Exception as e: - self._queue.put(_Result(value=None, error=e)) + except Exception as e: # noqa: BLE001 + self._queue.put(_Err(e)) self._loop.call_soon_threadsafe(self._yield_semaphore.release) finally: self._done_event.set() diff --git a/tests/unit/test_iterate.py b/tests/unit/test_iterate.py index 211e687..8e0e755 100644 --- a/tests/unit/test_iterate.py +++ b/tests/unit/test_iterate.py @@ -187,3 +187,37 @@ async def test_iterate_buffer_maxsize(streamcontext, buffer_maxsize: int | None) assert output == TEST_VALUES[:expected_len] await asyncio.wait_for(asyncio.get_running_loop().shutdown_default_executor(), 1.0) + + +class _ReciprocalIterator: + def __init__(self, iterator): + self._iter = iterator + + def __iter__(self): + return self + + def __next__(self): + return 1.0 / next(self._iter) + + +@athreading.iterate +def aiterate_recipricol(iterator): + return _ReciprocalIterator(iterator) + + +@pytest.mark.asyncio +async def test_iterate_background_exception(): + async with aiterate_recipricol(iter(range(-1, 2))) as stream: + with pytest.raises(ZeroDivisionError, match="float division by zero"): + [msg async for msg in stream] + + +@pytest.mark.xfail(reason="not supported yet") +@pytest.mark.asyncio +async def test_iterate_background_exception_suppress(): + async with aiterate_recipricol(iter(range(-1, 2))) as stream: + ait = stream.__aiter__() + assert await ait.__anext__() == -1 + with pytest.raises(ZeroDivisionError, match="float division by zero"): + await ait.__anext__() + assert await ait.__anext__() == 1