Propagate thread errors in iterator

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()

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