feat(request-loaders)!: Add get_pending_count to RequestLoader

docs/upgrading/upgrading_to_v2.md

@@ -343,6 +343,12 @@ class MyHttpClient(HttpClient):
 
 `HttpResponse` is now an abstract base class instead of a protocol, so a custom response subclasses it. It has three new abstract members, `url`, `headers_set_cookie` and `aclose`, and works as an async context manager that closes the response.
 
+## `get_pending_count` on request loaders
+
+`RequestLoader` has a new abstract `get_pending_count` method, which returns the number of requests that haven't been handled yet. If you have a custom request loader or manager, implement it. `RequestList` reads its source lazily, so like `get_total_count`, its `get_pending_count` leaves out the requests it hasn't fetched yet.
+
+`RequestManagerTandem.get_total_count` no longer counts a request twice once it moves from the loader to the manager. In v1, it added up the totals of both, so the crawler's status message showed too many pages in total.
+
 ## Logs changed
 
 If you configure loggers by name or match log messages, update them for the following changes.

src/crawlee/request_loaders/_request_list.py

@@ -159,6 +159,11 @@ async def get_handled_count(self) -> int:
     async def get_total_count(self) -> int:
         return self._assumed_total_count
 
+    @override
+    async def get_pending_count(self) -> int:
+        # The source is read lazily, so like the total count, this leaves out the requests not fetched yet.
+        return self._assumed_total_count - self._handled_count
+
     @override
     async def check_readiness(self) -> RequestLoaderStatus:
         await self._ensure_next_request()

src/crawlee/request_loaders/_request_loader.py

@@ -24,7 +24,7 @@ class RequestLoader(ABC):
     Key responsibilities:
         - Fetching the next request to be processed.
         - Marking requests as successfully handled after processing.
-        - Managing state information such as the total and handled request counts.
+        - Managing state information such as the total, pending and handled request counts.
     """
 
     @abstractmethod
@@ -35,6 +35,10 @@ async def get_handled_count(self) -> int:
     async def get_total_count(self) -> int:
         """Get an offline approximation of the total number of requests in the loader (i.e. pending + handled)."""
 
+    @abstractmethod
+    async def get_pending_count(self) -> int:
+        """Get an approximation of the number of requests in the loader that haven't been handled yet."""
+
     @abstractmethod
     async def check_readiness(self) -> RequestSourceStatus:
         """Report whether the loader has a request to hand over, is waiting for one, or is done.

src/crawlee/request_loaders/_request_manager_tandem.py

@@ -42,7 +42,14 @@ async def get_handled_count(self) -> int:
 
     @override
     async def get_total_count(self) -> int:
-        return (await self._read_only_loader.get_total_count()) + (await self._read_write_manager.get_total_count())
+        # A request moved into the manager counts as handled in the loader, so only the loader's pending ones are added.
+        loader_pending = await self._read_only_loader.get_pending_count()
+        return loader_pending + (await self._read_write_manager.get_total_count())
+
+    @override
+    async def get_pending_count(self) -> int:
+        loader_pending = await self._read_only_loader.get_pending_count()
+        return loader_pending + (await self._read_write_manager.get_pending_count())
 
     @override
     async def check_readiness(self) -> RequestSourceStatus:

src/crawlee/request_loaders/_sitemap_request_loader.py

@@ -209,6 +209,12 @@ async def get_handled_count(self) -> int:
         state = await self._get_state()
         return state.handled_count
 
+    @override
+    async def get_pending_count(self) -> int:
+        """Return the number of URLs found so far that haven't been handled yet."""
+        state = await self._get_state()
+        return len(state.url_queue) + len(state.in_progress)
+
     @override
     async def check_readiness(self) -> RequestLoaderStatus:
         state = await self._get_state()

src/crawlee/request_loaders/_throttling_request_manager.py

@@ -305,6 +305,15 @@ async def get_total_count(self) -> int:
         )
         return sum(counts) - self._migrated_from_inner
 
+    @override
+    async def get_pending_count(self) -> int:
+        await self._ensure_sub_managers()
+        counts = await asyncio.gather(
+            self._inner.get_pending_count(), *(sm.get_pending_count() for sm in self._sub_managers.values())
+        )
+        # `inner` counts a migrated request as handled, so it is pending only in its sub-manager and needs no fix-up.
+        return sum(counts)
+
     @override
     async def check_readiness(self) -> RequestSourceStatus:
         """Report whether anything can be dispatched now, and otherwise when it can be or why it never will.

src/crawlee/storages/_request_queue.py

@@ -116,6 +116,11 @@ async def get_total_count(self) -> int:
         metadata = await self._backend.get_metadata()
         return metadata.total_request_count
 
+    @override
+    async def get_pending_count(self) -> int:
+        metadata = await self._backend.get_metadata()
+        return metadata.total_request_count - metadata.handled_request_count
+
     @override
     @classmethod
     async def open(

tests/unit/request_loaders/test_request_list.py

@@ -36,6 +36,19 @@ async def generator() -> AsyncGenerator[str]:
     assert await request_list.get_handled_count() == 3
 
 
+async def test_pending_count_counts_fetched() -> None:
+    """Test that the pending count covers only the fetched requests not handled yet, since the source is lazy."""
+    request_list = RequestList(['https://a.placeholder.com', 'https://b.placeholder.com', 'https://c.placeholder.com'])
+
+    first = await request_list.fetch_next_request()
+    second = await request_list.fetch_next_request()
+    assert first is not None
+    assert second is not None
+    await request_list.mark_request_as_handled(first)
+
+    assert await request_list.get_pending_count() == 1
+
+
 async def test_readiness_during_traversal() -> None:
     """Test that the list waits for fetched requests and finishes once they're all handled."""
     request_list = RequestList(['https://a.placeholder.com', 'https://b.placeholder.com', 'https://c.placeholder.com'])

tests/unit/request_loaders/test_sitemap_request_loader.py

@@ -229,6 +229,23 @@ async def test_fetch_none_while_in_progress(server_url: URL, http_client: HttpCl
     assert (await sitemap_loader.check_readiness()).status == 'waiting'
 
 
+async def test_pending_count_during_traversal(server_url: URL, http_client: HttpClient) -> None:
+    """Test that the pending count covers the queued URLs and the ones in progress."""
+    sitemap_url = (server_url / 'sitemap.xml').with_query(
+        base64=encode_base64(get_basic_sitemap(url=server_url).encode())
+    )
+    sitemap_loader = SitemapRequestLoader([str(sitemap_url)], http_client=http_client, enqueue_strategy='all')
+
+    first = await sitemap_loader.fetch_next_request()
+    second = await sitemap_loader.fetch_next_request()
+    assert first is not None
+    assert second is not None
+    await sitemap_loader.mark_request_as_handled(first)
+
+    assert await poll_until_condition(sitemap_loader.get_total_count, lambda count: count == 5) == 5
+    assert await sitemap_loader.get_pending_count() == 4
+
+
 async def test_abort_sitemap_loading(server_url: URL, http_client: HttpClient) -> None:
     sitemap_url = (server_url / 'sitemap.xml').with_query(
         base64=encode_base64(get_basic_sitemap(url=server_url).encode())

tests/unit/storages/test_request_manager_tandem.py

@@ -10,6 +10,7 @@
 from crawlee.request_loaders import (
     FinishedStatus,
     ReadyStatus,
+    RequestList,
     RequestLoader,
     RequestManager,
     RequestManagerTandem,
@@ -157,3 +158,28 @@ async def test_transfers_only_when_loader_ready() -> None:
     assert request is not None
     assert request.url == 'https://a.placeholder.com'
     loader.fetch_next_request.assert_not_awaited()
+
+
+async def test_total_count_counts_moved_once() -> None:
+    """Test that a request moved from the loader to the manager counts once in the total."""
+    tandem = RequestManagerTandem(
+        RequestList(['https://a.placeholder.com', 'https://b.placeholder.com']), await RequestQueue.open()
+    )
+
+    while (await tandem.check_readiness()).status != 'finished':
+        request = await tandem.fetch_next_request()
+        assert request is not None
+        await tandem.mark_request_as_handled(request)
+
+    assert await tandem.get_total_count() == 2
+    assert await tandem.get_pending_count() == 0
+
+
+async def test_pending_count_sums_sources() -> None:
+    """Test that the pending count adds up the pending requests of the loader and the manager."""
+    loader = create_autospec(RequestLoader, instance=True, spec_set=True)
+    loader.get_pending_count.return_value = 2
+    manager = create_autospec(RequestManager, instance=True, spec_set=True)
+    manager.get_pending_count.return_value = 3
+
+    assert await RequestManagerTandem(loader, manager).get_pending_count() == 5

tests/unit/storages/test_request_queue.py

@@ -628,6 +628,17 @@ async def test_add_requests_wait_for_all(
     assert await rq.get_total_count() == 15
 
 
+async def test_get_pending_count(rq: RequestQueue) -> None:
+    """Test that the pending count leaves out the handled requests."""
+    await rq.add_requests(['https://example.com/1', 'https://example.com/2', 'https://example.com/3'])
+
+    request = await rq.fetch_next_request()
+    assert request is not None
+    await rq.mark_request_as_handled(request)
+
+    assert await rq.get_pending_count() == 2
+
+
 async def test_mark_non_existent_request_as_handled(rq: RequestQueue) -> None:
     """Test marking a non-existent request as handled."""
     # Create a request that hasn't been added to the queue

tests/unit/test_throttling_request_manager.py

@@ -891,6 +891,32 @@ async def test_get_total_count_aggregates(manager: ThrottlingRequestManager[Requ
     assert await manager.get_total_count() == 2
 
 
+async def test_get_pending_count_aggregates(manager: ThrottlingRequestManager[RequestQueue]) -> None:
+    """Test that the pending count sums inner and every sub-manager, without the handled requests."""
+    await manager.add_request(f'https://{THROTTLED_DOMAIN}/page1')
+    await manager.add_request(f'https://{NON_THROTTLED_DOMAIN}/page1')
+
+    request = await manager.fetch_next_request()
+    assert request is not None
+    await manager.mark_request_as_handled(request)
+
+    assert await manager.get_pending_count() == 1
+
+
+async def test_pending_count_after_migration(
+    manager: ThrottlingRequestManager[RequestQueue],
+    inner_queue: RequestQueue,
+) -> None:
+    """Test that a request migrated from inner to its domain's sub-manager counts once as pending."""
+    await inner_queue.add_requests([f'https://{THROTTLED_DOMAIN}/page{i}' for i in range(2)])
+    manager.record_domain_delay(f'https://{THROTTLED_DOMAIN}/', retry_after=timedelta(seconds=60))
+
+    assert await manager.fetch_next_request() is None
+
+    assert await inner_queue.get_handled_count() == 2
+    assert await manager.get_pending_count() == 2
+
+
 async def test_readiness_aggregates(manager: ThrottlingRequestManager[RequestQueue]) -> None:
     """Test that the readiness follows a request through the sub-manager it's routed to."""
     assert (await manager.check_readiness()).status == 'finished'