fix: Make result of RequestList.is_empty independent of fetch_next_request calls (#876)

janbuchar · web-flow · commit d50249ecbfe2 · 2025-01-07T14:16:40.000+01:00
The old version of `RequestList` only updated the `_is_empty` flag on
`fetch_next_request` calls, which was not enough. This PR updates it so
that the `is_empty` method also tries to dequeue a request from the
iterator before returning a result.
diff --git a/src/crawlee/request_loaders/_request_list.py b/src/crawlee/request_loaders/_request_list.py
@@ -39,13 +39,12 @@ def __init__(
         self._assumed_total_count = 0
 
         self._in_progress = set[str]()
-        self._is_empty = False
+        self._next: Request | None = None
 
         if isinstance(requests, AsyncIterable):
             self._requests = requests.__aiter__()
         elif requests is None:
             self._requests = self._iterate_in_threadpool([])
-            self._is_empty = True
         else:
             self._requests = self._iterate_in_threadpool(requests)
 
@@ -61,30 +60,27 @@ async def get_total_count(self) -> int:
 
     @override
     async def is_empty(self) -> bool:
-        return self._is_empty
+        await self._ensure_next_request()
+        return self._next is None
 
     @override
     async def is_finished(self) -> bool:
-        return self._is_empty and len(self._in_progress) == 0
+        return len(self._in_progress) == 0 and await self.is_empty()
 
     @override
     async def fetch_next_request(self) -> Request | None:
-        if self._is_empty:
+        await self._ensure_next_request()
+
+        if self._next is None:
             return None
 
-        if self._requests_lock is None:
-            self._requests_lock = asyncio.Lock()
+        self._in_progress.add(self._next.id)
+        self._assumed_total_count += 1
 
-        try:
-            async with self._requests_lock:
-                request = self._transform_request(await self._requests.__anext__())
-        except StopAsyncIteration:
-            self._is_empty = True
-            return None
-        else:
-            self._in_progress.add(request.id)
-            self._assumed_total_count += 1
-            return request
+        next_request = self._next
+        self._next = None
+
+        return next_request
 
     @override
     async def mark_request_as_handled(self, request: Request) -> None:
@@ -95,6 +91,17 @@ async def mark_request_as_handled(self, request: Request) -> None:
     async def get_handled_count(self) -> int:
         return self._handled_count
 
+    async def _ensure_next_request(self) -> None:
+        if self._requests_lock is None:
+            self._requests_lock = asyncio.Lock()
+
+        try:
+            async with self._requests_lock:
+                if self._next is None:
+                    self._next = self._transform_request(await self._requests.__anext__())
+        except StopAsyncIteration:
+            self._next = None
+
     async def _iterate_in_threadpool(self, iterable: Iterable[str | Request]) -> AsyncIterator[str | Request]:
         """Inspired by a function of the same name from encode/starlette."""
         iterator = iter(iterable)
@@ -115,4 +122,4 @@ def _next() -> str | Request:
             while True:
                 yield await asyncio.to_thread(_next)
         except _StopIteration:
-            raise StopAsyncIteration  # noqa: B904
+            return
diff --git a/tests/unit/request_loaders/test_request_list.py b/tests/unit/request_loaders/test_request_list.py
@@ -0,0 +1,58 @@
+from collections.abc import AsyncGenerator
+
+from crawlee.request_loaders._request_list import RequestList
+
+
+async def test_sync_traversal() -> None:
+    request_list = RequestList(['https://a.com', 'https://b.com', 'https://c.com'])
+
+    while not await request_list.is_finished():
+        item = await request_list.fetch_next_request()
+        assert item is not None
+
+        await request_list.mark_request_as_handled(item)
+
+    assert await request_list.is_empty()
+
+
+async def test_async_traversal() -> None:
+    async def generator() -> AsyncGenerator[str]:
+        yield 'https://a.com'
+        yield 'https://b.com'
+        yield 'https://c.com'
+
+    request_list = RequestList(generator())
+
+    while not await request_list.is_finished():
+        item = await request_list.fetch_next_request()
+        assert item is not None
+
+        await request_list.mark_request_as_handled(item)
+
+    assert await request_list.is_empty()
+
+
+async def test_is_empty_does_not_depend_on_fetch_next_request() -> None:
+    request_list = RequestList(['https://a.com', 'https://b.com', 'https://c.com'])
+
+    item_1 = await request_list.fetch_next_request()
+    assert item_1 is not None
+    assert not await request_list.is_finished()
+
+    item_2 = await request_list.fetch_next_request()
+    assert item_2 is not None
+    assert not await request_list.is_finished()
+
+    item_3 = await request_list.fetch_next_request()
+    assert item_3 is not None
+    assert not await request_list.is_finished()
+
+    assert await request_list.is_empty()
+    assert not await request_list.is_finished()
+
+    await request_list.mark_request_as_handled(item_1)
+    await request_list.mark_request_as_handled(item_2)
+    await request_list.mark_request_as_handled(item_3)
+
+    assert await request_list.is_empty()
+    assert await request_list.is_finished()
