diff --git a/docs/examples/code_examples/file_download.py b/docs/examples/code_examples/file_download.py new file mode 100644 index 0000000000..7a5cf65411 --- /dev/null +++ b/docs/examples/code_examples/file_download.py @@ -0,0 +1,43 @@ +import asyncio + +from yarl import URL + +from crawlee.crawlers import FileDownloadCrawler, FileDownloadCrawlingContext + + +async def main() -> None: + # FileDownloadCrawler downloads files with plain HTTP requests and accepts + # any content type. + crawler = FileDownloadCrawler( + # Limit the crawl to max requests. Remove or increase it for crawling all links. + max_requests_per_crawl=10, + ) + + # Define the default request handler, which will be called for every request. + @crawler.router.default_handler + async def request_handler(context: FileDownloadCrawlingContext) -> None: + context.log.info(f'Downloading {context.request.url} ...') + + # Read the whole file into memory. + content = await context.http_response.read() + + # Save the file to the default key-value store with the server's content type. + kvs = await context.get_key_value_store() + file_name = URL(context.request.url).name + await kvs.set_value( + key=file_name, + value=content, + content_type=context.http_response.headers.get('content-type'), + ) + + # Run the crawler with the list of files to download. + await crawler.run( + [ + 'https://pdfobject.com/pdf/sample.pdf', + 'https://crawlee.dev/assets/images/gradcracker-scraper-caefb62d1c150c4209a6e564c052fa41.webp', + ] + ) + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/docs/examples/code_examples/file_download_stream.py b/docs/examples/code_examples/file_download_stream.py new file mode 100644 index 0000000000..82f1b2018a --- /dev/null +++ b/docs/examples/code_examples/file_download_stream.py @@ -0,0 +1,49 @@ +import asyncio +from datetime import timedelta +from pathlib import Path + +from yarl import URL + +from crawlee.crawlers import FileDownloadCrawler, FileDownloadCrawlingContext + +DOWNLOAD_DIR = Path('downloads') + + +async def main() -> None: + # With stream=True, the request handler receives a response whose body has not + # been read yet. + crawler = FileDownloadCrawler( + stream=True, + # Bounds establishing the connection and receiving the response headers. + navigation_timeout=timedelta(minutes=5), + # The body is downloaded inside the handler, so this bounds the transfer itself. + request_handler_timeout=timedelta(minutes=5), + # Limit the crawl to max requests. Remove or increase it for crawling all links. + max_requests_per_crawl=10, + ) + + # Define the default request handler, which will be called for every request. + @crawler.router.default_handler + async def request_handler(context: FileDownloadCrawlingContext) -> None: + context.log.info(f'Downloading {context.request.url} ...') + + file_name = URL(context.request.url).name + + # Write each chunk to disk as it arrives, without buffering the whole file. + with (DOWNLOAD_DIR / file_name).open('wb') as file: + async for chunk in context.http_response.read_stream(): + file.write(chunk) + + context.log.info(f'Saved {file_name}') + + # Run the crawler with the list of files to download. + await crawler.run( + [ + 'https://samplelib.com/mp4/sample-15s-720p.mp4', + ] + ) + + +if __name__ == '__main__': + DOWNLOAD_DIR.mkdir(exist_ok=True) + asyncio.run(main()) diff --git a/docs/examples/file_download.mdx b/docs/examples/file_download.mdx new file mode 100644 index 0000000000..40c859cbe2 --- /dev/null +++ b/docs/examples/file_download.mdx @@ -0,0 +1,25 @@ +--- +id: file-download +title: Download files +--- + +import ApiLink from '@site/src/components/ApiLink'; +import RunnableCodeBlock from '@site/src/components/RunnableCodeBlock'; +import CodeBlock from '@theme/CodeBlock'; + +import FileDownloadExample from '!!raw-loader!roa-loader!./code_examples/file_download.py'; +import FileDownloadStreamExample from '!!raw-loader!./code_examples/file_download_stream.py'; + +This example demonstrates how to use `FileDownloadCrawler` to download files such as PDFs, images or videos with plain HTTP requests. The crawler doesn't parse the response body, so any content type is accepted. Each downloaded file is saved to the default `KeyValueStore` together with the content type reported by the server. + + + {FileDownloadExample} + + +## Streaming large files + +Buffering a whole file in memory doesn't scale to large downloads. Construct the crawler with `stream=True` and the request handler receives a response whose body hasn't been read yet. Consume it in chunks with `read_stream()` and write each chunk to disk as it arrives. + + + {FileDownloadStreamExample} + diff --git a/docs/guides/architecture_overview.mdx b/docs/guides/architecture_overview.mdx index 2eb1608fa1..c9f28a04d8 100644 --- a/docs/guides/architecture_overview.mdx +++ b/docs/guides/architecture_overview.mdx @@ -51,6 +51,8 @@ class BeautifulSoupCrawler class PydanticAiCrawler +class FileDownloadCrawler + class PlaywrightCrawler class AdaptivePlaywrightCrawler @@ -68,6 +70,7 @@ AbstractHttpCrawler --|> HttpCrawler AbstractHttpCrawler --|> ParselCrawler AbstractHttpCrawler --|> BeautifulSoupCrawler AbstractHttpCrawler --|> PydanticAiCrawler +AbstractHttpCrawler --|> FileDownloadCrawler PlaywrightCrawler --|> StagehandCrawler ``` @@ -75,12 +78,13 @@ PlaywrightCrawler --|> StagehandCrawler HTTP crawlers use HTTP clients to fetch pages and parse them with HTML parsing libraries. They are fast and efficient for sites that do not require JavaScript rendering. HTTP clients are Crawlee components that wrap around HTTP libraries like [httpx](https://www.python-httpx.org/), [curl-impersonate](https://github.com/lwthiker/curl-impersonate) or [impit](https://github.com/apify/impit) and handle HTTP communication for requests and responses. You can learn more about them in the [HTTP clients guide](./http-clients). -HTTP crawlers inherit from `AbstractHttpCrawler` and there are four crawlers that belong to this category: +HTTP crawlers inherit from `AbstractHttpCrawler` and there are five crawlers that belong to this category: - `BeautifulSoupCrawler` utilizes the [BeautifulSoup](https://www.crummy.com/software/BeautifulSoup/) HTML parser. - `ParselCrawler` utilizes [Parsel](https://github.com/scrapy/parsel) for parsing HTML. - `HttpCrawler` does not parse HTTP responses at all and is used when no content parsing is required. - `PydanticAiCrawler` parses HTML with Parsel and uses an LLM to extract structured data into a validated Pydantic model. +- `FileDownloadCrawler` downloads files of any content type, optionally streaming large bodies in chunks. You can learn more about HTTP crawlers in the [HTTP crawlers guide](./http-crawlers). @@ -126,6 +130,8 @@ class BeautifulSoupCrawlingContext class PydanticAiCrawlingContext +class FileDownloadCrawlingContext + class PlaywrightPreNavCrawlingContext class PlaywrightCrawlingContext @@ -156,6 +162,8 @@ ParsedHttpCrawlingContext --|> BeautifulSoupCrawlingContext ParselCrawlingContext --|> PydanticAiCrawlingContext +HttpCrawlingContext --|> FileDownloadCrawlingContext + BasicCrawlingContext --|> PlaywrightPreNavCrawlingContext PlaywrightPreNavCrawlingContext --|> PlaywrightCrawlingContext @@ -178,6 +186,7 @@ They have a similar inheritance structure as the crawlers, with the base class b - `ParselCrawlingContext` for HTTP crawlers that use [Parsel](https://github.com/scrapy/parsel) for parsing. - `BeautifulSoupCrawlingContext` for HTTP crawlers that use [BeautifulSoup](https://www.crummy.com/software/BeautifulSoup/) for parsing. - `PydanticAiCrawlingContext` for the AI crawler, extending the Parsel context with an `extract` helper. +- `FileDownloadCrawlingContext` for the file download crawler. - `PlaywrightPreNavCrawlingContext` for Playwright crawlers before the page is navigated. - `PlaywrightCrawlingContext` for Playwright crawlers. - `AdaptivePlaywrightPreNavCrawlingContext` for Adaptive Playwright crawlers before the page is navigated. diff --git a/docs/guides/http_crawlers.mdx b/docs/guides/http_crawlers.mdx index 366b36127c..00100d5e46 100644 --- a/docs/guides/http_crawlers.mdx +++ b/docs/guides/http_crawlers.mdx @@ -30,7 +30,7 @@ HTTP crawlers are ideal for extracting data from server-rendered websites that d ## Overview -All HTTP crawlers share a common architecture built around the `AbstractHttpCrawler` base class. The main differences lie in the parsing strategy and the context provided to request handlers. There are `BeautifulSoupCrawler`, `ParselCrawler`, and `HttpCrawler`. It can also be extended to create custom crawlers with specialized parsing requirements. They use HTTP clients to fetch page content and parsing libraries to extract data from the HTML, check out the [HTTP clients guide](./http-clients) to learn about the HTTP clients used by these crawlers, how to switch between them, and how to create custom HTTP clients tailored to your specific requirements. +All HTTP crawlers share a common architecture built around the `AbstractHttpCrawler` base class. The main differences lie in the parsing strategy and the context provided to request handlers. There are `BeautifulSoupCrawler`, `ParselCrawler`, `HttpCrawler`, and `FileDownloadCrawler`. It can also be extended to create custom crawlers with specialized parsing requirements. They use HTTP clients to fetch page content and parsing libraries to extract data from the HTML, check out the [HTTP clients guide](./http-clients) to learn about the HTTP clients used by these crawlers, how to switch between them, and how to create custom HTTP clients tailored to your specific requirements. ```mermaid --- @@ -63,6 +63,8 @@ class ParselCrawler class BeautifulSoupCrawler +class FileDownloadCrawler + %% ======================== %% Inheritance arrows %% ======================== @@ -71,6 +73,7 @@ BasicCrawler --|> AbstractHttpCrawler AbstractHttpCrawler --|> HttpCrawler AbstractHttpCrawler --|> ParselCrawler AbstractHttpCrawler --|> BeautifulSoupCrawler +AbstractHttpCrawler --|> FileDownloadCrawler ``` ## BeautifulSoupCrawler @@ -131,6 +134,10 @@ The following examples demonstrate how to integrate with several popular parsing +## FileDownloadCrawler + +The `FileDownloadCrawler` downloads files instead of scraping pages. It accepts any content type without parsing and gives the request handler direct access to the response body. By default the whole file is buffered in memory. For large files, construct the crawler with `stream=True` and consume the body in chunks via `read_stream()`. For usage, see the [Download files](../examples/file-download) example. + ## Custom HTTP crawler While the built-in crawlers cover most use cases, you might need a custom HTTP crawler for specialized parsing requirements. To create a custom HTTP crawler, inherit directly from `AbstractHttpCrawler`. This approach requires implementing: diff --git a/src/crawlee/crawlers/__init__.py b/src/crawlee/crawlers/__init__.py index 4bc9827db3..e784c2c8ee 100644 --- a/src/crawlee/crawlers/__init__.py +++ b/src/crawlee/crawlers/__init__.py @@ -3,6 +3,7 @@ from ._abstract_http import AbstractHttpCrawler, AbstractHttpParser, HttpCrawlerOptions, ParsedHttpCrawlingContext from ._basic import BasicCrawler, BasicCrawlerOptions, BasicCrawlingContext, ContextPipeline +from ._file_download import FileDownloadCrawler, FileDownloadCrawlingContext from ._http import HttpCrawler, HttpCrawlingContext, HttpCrawlingResult _install_import_hook(__name__) @@ -113,6 +114,8 @@ 'BeautifulSoupCrawlingContext', 'BeautifulSoupParserType', 'ContextPipeline', + 'FileDownloadCrawler', + 'FileDownloadCrawlingContext', 'HttpCrawler', 'HttpCrawlerOptions', 'HttpCrawlingContext', diff --git a/src/crawlee/crawlers/_abstract_http/_abstract_http_crawler.py b/src/crawlee/crawlers/_abstract_http/_abstract_http_crawler.py index 8d15a1d801..fb85d71d4d 100644 --- a/src/crawlee/crawlers/_abstract_http/_abstract_http_crawler.py +++ b/src/crawlee/crawlers/_abstract_http/_abstract_http_crawler.py @@ -30,7 +30,7 @@ from ._abstract_http_parser import AbstractHttpParser -TCrawlingContext = TypeVar('TCrawlingContext', bound=ParsedHttpCrawlingContext) +TCrawlingContext = TypeVar('TCrawlingContext', bound=HttpCrawlingContext) TStatisticsState = TypeVar('TStatisticsState', bound=StatisticsState, default=StatisticsState) diff --git a/src/crawlee/crawlers/_file_download/__init__.py b/src/crawlee/crawlers/_file_download/__init__.py new file mode 100644 index 0000000000..80bd36056b --- /dev/null +++ b/src/crawlee/crawlers/_file_download/__init__.py @@ -0,0 +1,7 @@ +from ._file_download_crawler import FileDownloadCrawler +from ._file_download_crawling_context import FileDownloadCrawlingContext + +__all__ = [ + 'FileDownloadCrawler', + 'FileDownloadCrawlingContext', +] diff --git a/src/crawlee/crawlers/_file_download/_file_download_crawler.py b/src/crawlee/crawlers/_file_download/_file_download_crawler.py new file mode 100644 index 0000000000..0d05c6eba9 --- /dev/null +++ b/src/crawlee/crawlers/_file_download/_file_download_crawler.py @@ -0,0 +1,158 @@ +from __future__ import annotations + +from contextlib import AsyncExitStack +from typing import TYPE_CHECKING + +from crawlee._request import RequestState +from crawlee._utils.docs import docs_group +from crawlee.crawlers._abstract_http import AbstractHttpCrawler +from crawlee.crawlers._abstract_http._http_crawling_context import HttpCrawlingContext +from crawlee.crawlers._basic import ContextPipeline +from crawlee.crawlers._http._http_parser import NoParser + +from ._file_download_crawling_context import FileDownloadCrawlingContext + +if TYPE_CHECKING: + from collections.abc import AsyncGenerator + + from typing_extensions import Unpack + + from crawlee._types import BasicCrawlingContext + from crawlee.crawlers._abstract_http import HttpCrawlerOptions + + +@docs_group('Crawlers') +class FileDownloadCrawler(AbstractHttpCrawler[FileDownloadCrawlingContext, bytes, bytes]): + """A crawler for downloading files with plain HTTP requests. + + The `FileDownloadCrawler` builds on top of the `AbstractHttpCrawler`, which means it inherits all of its + features. Unlike the other HTTP crawlers, it does not parse the response body. Any content type is accepted + as-is, which makes it suitable for downloading binary files such as PDFs, images, videos or archives. + + By default the whole response body is buffered in memory and available via + `await context.http_response.read()`. For large files, construct the crawler with `stream=True`. The request + handler then receives a response whose body has not been read yet and can consume it in chunks via + `context.http_response.read_stream()`. The connection stays open while the handler runs and is closed + afterwards. A handler that returns without consuming the body just skips the download. + + ### Usage + + ```python + from yarl import URL + + from crawlee.crawlers import FileDownloadCrawler, FileDownloadCrawlingContext + + crawler = FileDownloadCrawler() + + # Define the default request handler, which will be called for every request. + @crawler.router.default_handler + async def request_handler(context: FileDownloadCrawlingContext) -> None: + context.log.info(f'Downloading {context.request.url} ...') + + # Read the whole file and save it to the key-value store. + content = await context.http_response.read() + kvs = await context.get_key_value_store() + await kvs.set_value( + key=URL(context.request.url).name, + value=content, + content_type=context.http_response.headers.get('content-type'), + ) + + await crawler.run(['https://example.com/document.pdf']) + ``` + + With `stream=True`, the body can be written to a file in chunks without loading it into memory: + + ```python + from pathlib import Path + + from yarl import URL + + from crawlee.crawlers import FileDownloadCrawler, FileDownloadCrawlingContext + + crawler = FileDownloadCrawler(stream=True) + + @crawler.router.default_handler + async def request_handler(context: FileDownloadCrawlingContext) -> None: + file_name = URL(context.request.url).name + + # Write chunks to disk as they arrive. + with Path(file_name).open('wb') as file: + async for chunk in context.http_response.read_stream(): + file.write(chunk) + + await crawler.run(['https://example.com/large-video.mp4']) + ``` + """ + + def __init__( + self, + *, + stream: bool = False, + **kwargs: Unpack[HttpCrawlerOptions[FileDownloadCrawlingContext]], + ) -> None: + """Initialize a new instance. + + Args: + stream: Whether to stream response bodies instead of buffering them in memory. In stream mode + the request handler consumes the body via `context.http_response.read_stream()`. + kwargs: Additional keyword arguments to pass to the underlying `AbstractHttpCrawler`. + """ + self._stream_response_body = stream + + kwargs['_context_pipeline'] = self._create_file_download_pipeline() + + super().__init__( + parser=NoParser(), + **kwargs, + ) + + def _create_file_download_pipeline(self) -> ContextPipeline[FileDownloadCrawlingContext]: + """Create the file download context pipeline with expected pipeline steps.""" + make_http_request = self._make_http_stream_request if self._stream_response_body else self._make_http_request + + return ( + ContextPipeline() + .compose(self._manage_shared_navigation_timeout) + .compose(self._execute_pre_navigation_hooks) + .compose(make_http_request) + .compose(self._execute_post_navigation_hooks) + .compose(self._handle_status_code_response) + .compose(self._to_file_download_crawling_context) + ) + + async def _make_http_stream_request( + self, context: BasicCrawlingContext + ) -> AsyncGenerator[HttpCrawlingContext, None]: + """Make a streamed HTTP request and create context enhanced by the streamed HTTP response. + + Args: + context: The current crawling context. + + Yields: + The original crawling context enhanced by the streamed HTTP response. + """ + async with AsyncExitStack() as exit_stack: + async with self._shared_navigation_timeouts[id(context.request)] as remaining_timeout: + response = await exit_stack.enter_async_context( + self._http_client.stream( + url=context.request.url, + method=context.request.method, + headers=context.request.headers, + payload=context.request.payload, + session=context.session, + proxy_info=context.proxy_info, + timeout=remaining_timeout, + ) + ) + + self._statistics.register_status_code(response.status_code) + context.request.state = RequestState.AFTER_NAV + + yield HttpCrawlingContext.from_basic_crawling_context(context=context, http_response=response) + + async def _to_file_download_crawling_context( + self, context: HttpCrawlingContext + ) -> AsyncGenerator[FileDownloadCrawlingContext, None]: + """Convert the HTTP crawling context to the final file download crawling context.""" + yield FileDownloadCrawlingContext.from_http_crawling_context(context) diff --git a/src/crawlee/crawlers/_file_download/_file_download_crawling_context.py b/src/crawlee/crawlers/_file_download/_file_download_crawling_context.py new file mode 100644 index 0000000000..9d17d4a084 --- /dev/null +++ b/src/crawlee/crawlers/_file_download/_file_download_crawling_context.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +from dataclasses import dataclass, fields + +from typing_extensions import Self, override + +from crawlee._types import PageSnapshot +from crawlee._utils.docs import docs_group +from crawlee.crawlers._abstract_http._http_crawling_context import HttpCrawlingContext + + +@dataclass(frozen=True) +@docs_group('Crawling contexts') +class FileDownloadCrawlingContext(HttpCrawlingContext): + """The crawling context used by the `FileDownloadCrawler`.""" + + @classmethod + def from_http_crawling_context(cls, context: HttpCrawlingContext) -> Self: + """Initialize a new instance from an existing `HttpCrawlingContext`.""" + return cls(**{field.name: getattr(context, field.name) for field in fields(context)}) + + @override + async def get_snapshot(self) -> PageSnapshot: + # Downloaded files are binary and streamed bodies cannot be re-read, so there is no page to snapshot. + return PageSnapshot() diff --git a/tests/unit/crawlers/_file_download/test_file_download.py b/tests/unit/crawlers/_file_download/test_file_download.py new file mode 100644 index 0000000000..02d9fffc9e --- /dev/null +++ b/tests/unit/crawlers/_file_download/test_file_download.py @@ -0,0 +1,203 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +from unittest.mock import AsyncMock, Mock, call + +import pytest + +from crawlee.crawlers import FileDownloadCrawler, FileDownloadCrawlingContext +from tests.unit.server import generate_file_content + +if TYPE_CHECKING: + from yarl import URL + + from crawlee.crawlers import HttpCrawlingContext + from crawlee.http_clients._base import HttpClient + +FILE_SIZE = 256 * 1024 +CHUNK_SIZE = 16 * 1024 + + +async def test_read_file(server_url: URL, http_client: HttpClient) -> None: + """The whole file is read into memory.""" + crawler = FileDownloadCrawler(http_client=http_client) + handler_call = Mock() + + @crawler.router.default_handler + async def request_handler(context: FileDownloadCrawlingContext) -> None: + handler_call( + content_type=context.http_response.headers.get('content-type'), + body=await context.http_response.read(), + ) + + stats = await crawler.run([str(server_url.with_path('file').with_query(size=FILE_SIZE))]) + + assert stats.requests_failed == 0 + assert stats.requests_finished == 1 + handler_call.assert_called_once() + assert handler_call.call_args.kwargs['content_type'] == 'application/octet-stream' + assert handler_call.call_args.kwargs['body'] == generate_file_content(FILE_SIZE) + + +async def test_stream_file(server_url: URL, http_client: HttpClient) -> None: + """The body is read in chunks, with metadata available before consumption.""" + crawler = FileDownloadCrawler(stream=True, http_client=http_client) + handler_call = Mock() + + @crawler.router.default_handler + async def request_handler(context: FileDownloadCrawlingContext) -> None: + # Metadata is passed before the body argument, so it is read before the stream is consumed. + handler_call( + status_code=context.http_response.status_code, + content_type=context.http_response.headers.get('content-type'), + content_length=context.http_response.headers.get('content-length'), + body=b''.join([chunk async for chunk in context.http_response.read_stream()]), + ) + + stats = await crawler.run([str(server_url.with_path('file').with_query(size=FILE_SIZE, chunk_size=CHUNK_SIZE))]) + + assert stats.requests_failed == 0 + assert stats.requests_finished == 1 + handler_call.assert_called_once() + assert handler_call.call_args.kwargs['status_code'] == 200 + assert handler_call.call_args.kwargs['content_type'] == 'application/octet-stream' + assert handler_call.call_args.kwargs['content_length'] == str(FILE_SIZE) + assert handler_call.call_args.kwargs['body'] == generate_file_content(FILE_SIZE) + + +async def test_slow_download(server_url: URL) -> None: + """Chunked responses with pauses between chunks download fully.""" + crawler = FileDownloadCrawler(stream=True) + handler_call = Mock() + + @crawler.router.default_handler + async def request_handler(context: FileDownloadCrawlingContext) -> None: + handler_call(body=b''.join([chunk async for chunk in context.http_response.read_stream()])) + + url = server_url.with_path('file').with_query(size=8 * 1024, chunk_size=1024, throttle=0.2) + stats = await crawler.run([str(url)]) + + assert stats.requests_failed == 0 + handler_call.assert_called_once() + assert handler_call.call_args.kwargs['body'] == generate_file_content(8 * 1024) + + +async def test_unconsumed_body(server_url: URL, http_client: HttpClient) -> None: + """A handler that never reads the body finishes cleanly instead of hanging.""" + crawler = FileDownloadCrawler(stream=True, http_client=http_client) + handler_call = Mock() + + @crawler.router.default_handler + async def request_handler(context: FileDownloadCrawlingContext) -> None: + handler_call(status_code=context.http_response.status_code) + + stats = await crawler.run([str(server_url.with_path('file').with_query(size=FILE_SIZE))]) + + assert stats.requests_failed == 0 + assert stats.requests_finished == 1 + handler_call.assert_called_once() + assert handler_call.call_args.kwargs['status_code'] == 200 + + +async def test_partially_consumed_body(server_url: URL, http_client: HttpClient) -> None: + """A handler that stops reading early finishes without a retry.""" + crawler = FileDownloadCrawler(stream=True, http_client=http_client) + handler_call = Mock() + + @crawler.router.default_handler + async def request_handler(context: FileDownloadCrawlingContext) -> None: + handler_call() + async for _ in context.http_response.read_stream(): + break + + stats = await crawler.run([str(server_url.with_path('file').with_query(size=FILE_SIZE, chunk_size=CHUNK_SIZE))]) + + assert stats.requests_failed == 0 + assert stats.requests_finished == 1 + handler_call.assert_called_once() + + +async def test_retry_failed_download(server_url: URL, http_client: HttpClient) -> None: + """A mid-stream failure retries and gets a fresh, complete response.""" + crawler = FileDownloadCrawler(stream=True, http_client=http_client, max_request_retries=1) + handler_call = Mock() + + @crawler.router.default_handler + async def request_handler(context: FileDownloadCrawlingContext) -> None: + stream = context.http_response.read_stream() + # Fail mid-stream on the first attempt, then read the full body on the retry. + if handler_call.call_count == 0: + await anext(stream) + handler_call(body=None) + raise RuntimeError('Simulated handler failure') + handler_call(body=b''.join([chunk async for chunk in stream])) + + stats = await crawler.run([str(server_url.with_path('file').with_query(size=FILE_SIZE, chunk_size=CHUNK_SIZE))]) + + assert stats.requests_failed == 0 + assert stats.requests_finished == 1 + assert handler_call.call_count == 2 + assert handler_call.call_args.kwargs['body'] == generate_file_content(FILE_SIZE) + + +@pytest.mark.parametrize('stream', [False, True], ids=['buffered', 'stream']) +async def test_error_status_code(server_url: URL, *, stream: bool) -> None: + """An error status code fails the request without calling the handler.""" + handler = AsyncMock() + crawler = FileDownloadCrawler(stream=stream, max_request_retries=0, request_handler=handler) + + stats = await crawler.run([str(server_url / 'status/404')]) + + assert stats.requests_failed == 1 + handler.assert_not_called() + + +@pytest.mark.parametrize('content_type', ['application/pdf', 'image/png', 'video/mp4']) +async def test_any_content_type(server_url: URL, content_type: str) -> None: + """Any content type is downloaded as-is.""" + crawler = FileDownloadCrawler() + handler_call = Mock() + + @crawler.router.default_handler + async def request_handler(context: FileDownloadCrawlingContext) -> None: + handler_call( + content_type=context.http_response.headers.get('content-type'), + body=await context.http_response.read(), + ) + + stats = await crawler.run([str(server_url.with_path('file').with_query(size=1024, content_type=content_type))]) + + assert stats.requests_failed == 0 + handler_call.assert_called_once() + assert handler_call.call_args.kwargs['content_type'] == content_type + assert handler_call.call_args.kwargs['body'] == generate_file_content(1024) + + +async def test_navigation_hooks(server_url: URL) -> None: + """Navigation hooks run around a streamed request, before the handler.""" + crawler = FileDownloadCrawler(stream=True) + tracker = Mock() + + @crawler.pre_navigation_hook + async def pre_nav_hook(_context: object) -> None: + tracker('pre_navigation') + + @crawler.post_navigation_hook + async def post_nav_hook(context: HttpCrawlingContext) -> None: + # The streamed response is already open when post-navigation hooks run. + tracker('post_navigation', context.http_response.status_code) + + @crawler.router.default_handler + async def request_handler(context: FileDownloadCrawlingContext) -> None: + tracker('handler') + async for _ in context.http_response.read_stream(): + pass + + stats = await crawler.run([str(server_url.with_path('file').with_query(size=1024))]) + + assert stats.requests_failed == 0 + assert tracker.mock_calls == [ + call('pre_navigation'), + call('post_navigation', 200), + call('handler'), + ]