102 lines
3.6 KiB
Python
102 lines
3.6 KiB
Python
import asyncio
|
|
from collections.abc import Callable
|
|
|
|
from yarl import URL
|
|
|
|
from crawlee import RequestOptions, RequestTransformAction
|
|
from crawlee.crawlers import BeautifulSoupCrawler, BeautifulSoupCrawlingContext
|
|
from crawlee.http_clients import ImpitHttpClient
|
|
from crawlee.request_loaders import SitemapRequestLoader
|
|
|
|
|
|
# Create a transform_request_function that maps request options based on the host in
|
|
# the URL
|
|
def create_transform_request(
|
|
data_mapper: dict[str, dict],
|
|
) -> Callable[[RequestOptions], RequestOptions | RequestTransformAction]:
|
|
def transform_request(
|
|
request_options: RequestOptions,
|
|
) -> RequestOptions | RequestTransformAction:
|
|
# According to the Sitemap protocol, all URLs in a Sitemap must be from a single
|
|
# host.
|
|
request_host = URL(request_options['url']).host
|
|
|
|
if request_host and (mapping_data := data_mapper.get(request_host)):
|
|
# Set properties from the mapping data
|
|
if 'label' in mapping_data:
|
|
request_options['label'] = mapping_data['label']
|
|
if 'user_data' in mapping_data:
|
|
request_options['user_data'] = mapping_data['user_data']
|
|
|
|
return request_options
|
|
|
|
return 'unchanged'
|
|
|
|
return transform_request
|
|
|
|
|
|
async def main() -> None:
|
|
# Prepare data mapping for hosts
|
|
apify_host = URL('https://apify.com/sitemap.xml').host
|
|
crawlee_host = URL('https://crawlee.dev/sitemap.xml').host
|
|
|
|
if not apify_host or not crawlee_host:
|
|
raise ValueError('Unable to extract host from URLs')
|
|
|
|
data_map = {
|
|
apify_host: {
|
|
'label': 'apify',
|
|
'user_data': {'source': 'apify'},
|
|
},
|
|
crawlee_host: {
|
|
'label': 'crawlee',
|
|
'user_data': {'source': 'crawlee'},
|
|
},
|
|
}
|
|
|
|
# Initialize the SitemapRequestLoader with the transform function
|
|
async with SitemapRequestLoader(
|
|
# Set the sitemap URLs and the HTTP client
|
|
sitemap_urls=['https://crawlee.dev/sitemap.xml', 'https://apify.com/sitemap.xml'],
|
|
http_client=ImpitHttpClient(),
|
|
transform_request_function=create_transform_request(data_map),
|
|
) as sitemap_loader:
|
|
# Convert the sitemap loader to a request manager
|
|
request_manager = await sitemap_loader.to_tandem()
|
|
|
|
# Create and configure the crawler
|
|
crawler = BeautifulSoupCrawler(
|
|
request_manager=request_manager,
|
|
max_requests_per_crawl=10,
|
|
)
|
|
|
|
# Create default handler for requests without a specific label
|
|
@crawler.router.default_handler
|
|
async def handler(context: BeautifulSoupCrawlingContext) -> None:
|
|
source = context.request.user_data.get('source', 'unknown')
|
|
context.log.info(
|
|
f'Processing request: {context.request.url} from source: {source}'
|
|
)
|
|
|
|
# Create handler for requests labeled 'apify'
|
|
@crawler.router.handler('apify')
|
|
async def apify_handler(context: BeautifulSoupCrawlingContext) -> None:
|
|
source = context.request.user_data.get('source', 'unknown')
|
|
context.log.info(
|
|
f'Apify handler processing: {context.request.url} from source: {source}'
|
|
)
|
|
|
|
# Create handler for requests labeled 'crawlee'
|
|
@crawler.router.handler('crawlee')
|
|
async def crawlee_handler(context: BeautifulSoupCrawlingContext) -> None:
|
|
source = context.request.user_data.get('source', 'unknown')
|
|
context.log.info(
|
|
f'Crawlee handler processing: {context.request.url} from source: {source}'
|
|
)
|
|
|
|
await crawler.run()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
asyncio.run(main())
|