Draft of use_state as input argument

Only for discussion, types ignored for now.

Author: Pijukatel

src/crawlee/crawlers/_adaptive_playwright/_adaptive_playwright_crawler.py

@@ -43,8 +43,7 @@
 
     from typing_extensions import Unpack
 
-    from crawlee.crawlers._basic._basic_crawler import _BasicCrawlerOptions
-
+from crawlee.crawlers._basic._basic_crawler import _BasicCrawlerOptions, _DefaultUseState
 
 TStaticParseResult = TypeVar('TStaticParseResult')
 TStaticSelectResult = TypeVar('TStaticSelectResult')
@@ -389,7 +388,8 @@ async def _run_request_handler(self, context: BasicCrawlingContext) -> None:
             # (This static crawl is performed only to evaluate rendering type detection.)
             kvs = await context.get_key_value_store()
             default_value = dict[str, JsonSerializable]()
-            old_state: dict[str, JsonSerializable] = await kvs.get_value(self._CRAWLEE_STATE_KEY, default_value)
+            # This was fragile even before. Out of scope for draft
+            old_state: dict[str, JsonSerializable] = await kvs.get_value(_DefaultUseState._CRAWLEE_STATE_KEY, default_value)
             old_state_copy = deepcopy(old_state)
 
         pw_run = await self._crawl_one('client only', context=context)

src/crawlee/crawlers/_basic/_basic_crawler.py

@@ -10,7 +10,7 @@
 import threading
 import traceback
 from asyncio import CancelledError
-from collections.abc import AsyncGenerator, Awaitable, Callable, Iterable, Sequence
+from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Iterable, Sequence
 from contextlib import AsyncExitStack, suppress
 from datetime import timedelta
 from functools import partial
@@ -42,6 +42,7 @@
     RequestHandlerRunResult,
     SendRequestFunction,
     SkippedReason,
+    UseStateFunction,
 )
 from crawlee._utils.docs import docs_group
 from crawlee._utils.file import atomic_write, export_csv_to_stream, export_json_to_stream
@@ -239,6 +240,29 @@ class BasicCrawlerOptions(
     """
 
 
+class _DefaultUseState:
+    _next_state_id = 0
+    _CRAWLEE_STATE_KEY = 'CRAWLEE_STATE'
+
+    def __init__(self, get_key_value_store: Awaitable[KeyValueStore]) -> None:
+        self._get_key_value_store = get_key_value_store
+        self._id = self._next_state_id
+        _DefaultUseState._next_state_id += 1
+
+    async def _use_state(
+        self,
+        default_value: dict[str, JsonSerializable] | None = None,
+    ) -> dict[str, JsonSerializable]:
+        kvs = await self._get_key_value_store()
+        return await kvs.get_auto_saved_value(f'{self._CRAWLEE_STATE_KEY}_{self._id}', default_value)
+
+    def __call__(
+        self,
+        default_value: dict[str, JsonSerializable] | None = None,
+    ) -> Coroutine[None, None, dict[str, JsonSerializable]]:
+        return self._use_state(default_value)
+
+
 @docs_group('Crawlers')
 class BasicCrawler(Generic[TCrawlingContext, TStatisticsState]):
     """A basic web crawler providing a framework for crawling websites.
@@ -264,7 +288,6 @@ class BasicCrawler(Generic[TCrawlingContext, TStatisticsState]):
         - and more.
     """
 
-    _CRAWLEE_STATE_KEY = 'CRAWLEE_STATE'
     _request_handler_timeout_text = 'Request handler timed out after'
     __next_id = 0
 
@@ -298,7 +321,7 @@ def __init__(
         status_message_logging_interval: timedelta = timedelta(seconds=10),
         status_message_callback: Callable[[StatisticsState, StatisticsState | None, str], Awaitable[str | None]]
         | None = None,
-        crawler_id: int | None = None,
+        use_state: UseStateFunction | None = None,
         _context_pipeline: ContextPipeline[TCrawlingContext] | None = None,
         _additional_context_managers: Sequence[AbstractAsyncContextManager] | None = None,
         _logger: logging.Logger | None = None,
@@ -351,22 +374,15 @@ def __init__(
             status_message_logging_interval: Interval for logging the crawler status messages.
             status_message_callback: Allows overriding the default status message. The default status message is
                 provided in the parameters. Returning `None` suppresses the status message.
-            crawler_id: Id of the crawler used for state and statistics tracking. You can use same explicit id to share
-                state and statistics between two crawlers. By default, each crawler will use own state and statistics.
+            use_state: Callback used to access shared state. Use only for custom state implementation, for example when
+                you want to share state between two different crawlers.
             _context_pipeline: Enables extending the request lifecycle and modifying the crawling context.
                 Intended for use by subclasses rather than direct instantiation of `BasicCrawler`.
             _additional_context_managers: Additional context managers used throughout the crawler lifecycle.
                 Intended for use by subclasses rather than direct instantiation of `BasicCrawler`.
             _logger: A logger instance, typically provided by a subclass, for consistent logging labels.
                 Intended for use by subclasses rather than direct instantiation of `BasicCrawler`.
         """
-        if crawler_id is None:
-            # This could look into set of already used ids, but lets not overengineer this.
-            self.id = BasicCrawler.__next_id
-            BasicCrawler.__next_id += 1
-        else:
-            self.id = crawler_id
-
         implicit_event_manager_with_explicit_config = False
         if not configuration:
             configuration = service_locator.get_configuration()
@@ -442,6 +458,12 @@ def __init__(
         self._use_session_pool = use_session_pool
         self._retry_on_blocked = retry_on_blocked
 
+        # Set use state
+        if use_state:
+            self._use_state = use_state
+        else:
+            self._use_state = _DefaultUseState(get_key_value_store=self.get_key_value_store)
+
         # Logging setup
         if configure_logging:
             root_logger = logging.getLogger()
@@ -837,13 +859,6 @@ async def add_requests(
             wait_for_all_requests_to_be_added_timeout=wait_for_all_requests_to_be_added_timeout,
         )
 
-    async def _use_state(
-        self,
-        default_value: dict[str, JsonSerializable] | None = None,
-    ) -> dict[str, JsonSerializable]:
-        kvs = await self.get_key_value_store()
-        return await kvs.get_auto_saved_value(f'{self._CRAWLEE_STATE_KEY}_{self.id}', default_value)
-
     async def _save_crawler_state(self) -> None:
         store = await self.get_key_value_store()
         await store.persist_autosaved_values()

tests/unit/conftest.py

@@ -12,7 +12,7 @@
 from uvicorn.config import Config
 
 from crawlee import service_locator
-from crawlee.crawlers import BasicCrawler
+from crawlee.crawlers._basic._basic_crawler import _DefaultUseState
 from crawlee.fingerprint_suite._browserforge_adapter import get_available_header_network
 from crawlee.http_clients import CurlImpersonateHttpClient, HttpxHttpClient, ImpitHttpClient
 from crawlee.proxy_configuration import ProxyInfo
@@ -75,7 +75,7 @@ def _prepare_test_env() -> None:
         # Reset global class variables to ensure test isolation.
         KeyValueStore._autosaved_values = {}
         Statistics._Statistics__next_id = 0  # type:ignore[attr-defined] # Mangled attribute
-        BasicCrawler._BasicCrawler__next_id = 0  # type:ignore[attr-defined] # Mangled attribute
+        _DefaultUseState._next_state_id = 0  # type:ignore[attr-defined] # Mangled attribute
 
     return _prepare_test_env
 

tests/unit/crawlers/_adaptive_playwright/test_adaptive_playwright_crawler.py

@@ -19,7 +19,6 @@
     AdaptivePlaywrightCrawler,
     AdaptivePlaywrightCrawlingContext,
     AdaptivePlaywrightPreNavCrawlingContext,
-    BasicCrawler,
     RenderingType,
     RenderingTypePrediction,
     RenderingTypePredictor,
@@ -381,7 +380,7 @@ async def test_adaptive_crawling_result_use_state_isolation(
     crawler = AdaptivePlaywrightCrawler.with_beautifulsoup_static_parser(
         rendering_type_predictor=static_only_predictor_enforce_detection,
     )
-    await key_value_store.set_value(f'{BasicCrawler._CRAWLEE_STATE_KEY}_0', {'counter': 0})
+    await key_value_store.set_value('CRAWLEE_STATE_0', {'counter': 0})
     request_handler_calls = 0
 
     @crawler.router.default_handler
@@ -398,7 +397,7 @@ async def request_handler(context: AdaptivePlaywrightCrawlingContext) -> None:
     # Request handler was called twice
     assert request_handler_calls == 2
     # Increment of global state happened only once
-    assert (await key_value_store.get_value(f'{BasicCrawler._CRAWLEE_STATE_KEY}_0'))['counter'] == 1
+    assert (await key_value_store.get_value('CRAWLEE_STATE_0'))['counter'] == 1
 
 
 async def test_adaptive_crawling_statistics(test_urls: list[str]) -> None:

tests/unit/crawlers/_basic/test_basic_crawler.py

@@ -810,36 +810,39 @@ async def handler(context: BasicCrawlingContext) -> None:
     await crawler.run(['https://hello.world'])
 
     kvs = await crawler.get_key_value_store()
-    value = await kvs.get_value(f'{BasicCrawler._CRAWLEE_STATE_KEY}_0')
+    value = await kvs.get_value('CRAWLEE_STATE_0')
 
     assert value == {'hello': 'world'}
 
 
-async def test_context_use_state_crawlers_share_state() -> None:
+async def test_context_use_state_crawlers_share_custom_state() -> None:
+    custom_state_dict = {}
+
+    async def custom_use_state(default_state: dict[str, JsonSerializable]) -> dict[str, JsonSerializable]:
+        if not custom_state_dict:
+            custom_state_dict.update(default_state)
+        return custom_state_dict
+
     async def handler(context: BasicCrawlingContext) -> None:
         state = await context.use_state({'urls': []})
         assert isinstance(state['urls'], list)
         state['urls'].append(context.request.url)
 
-    crawler_1 = BasicCrawler(crawler_id=0, request_handler=handler)
-    crawler_2 = BasicCrawler(crawler_id=0, request_handler=handler)
+    crawler_1 = BasicCrawler(use_state=custom_use_state, request_handler=handler)
+    crawler_2 = BasicCrawler(use_state=custom_use_state, request_handler=handler)
 
     await crawler_1.run(['https://a.com'])
     await crawler_2.run(['https://b.com'])
 
-    kvs = await KeyValueStore.open()
-    assert crawler_1.id == crawler_2.id == 0
-    assert await kvs.get_value(f'{BasicCrawler._CRAWLEE_STATE_KEY}_{crawler_1.id}') == {
-        'urls': ['https://a.com', 'https://b.com']
-    }
+    assert custom_state_dict == {'urls': ['https://a.com', 'https://b.com']}
 
 
 async def test_crawlers_share_stats() -> None:
     async def handler(context: BasicCrawlingContext) -> None:
         await context.use_state({'urls': []})
 
-    crawler_1 = BasicCrawler(crawler_id=0, request_handler=handler)
-    crawler_2 = BasicCrawler(crawler_id=0, request_handler=handler, statistics=crawler_1.statistics)
+    crawler_1 = BasicCrawler(request_handler=handler)
+    crawler_2 = BasicCrawler(request_handler=handler, statistics=crawler_1.statistics)
 
     result1 = await crawler_1.run(['https://a.com'])
     result2 = await crawler_2.run(['https://b.com'])
@@ -862,8 +865,8 @@ async def handler(context: BasicCrawlingContext) -> None:
     await crawler_2.run(['https://b.com'])
 
     kvs = await KeyValueStore.open()
-    assert await kvs.get_value(f'{BasicCrawler._CRAWLEE_STATE_KEY}_0') == {'urls': ['https://a.com']}
-    assert await kvs.get_value(f'{BasicCrawler._CRAWLEE_STATE_KEY}_1') == {'urls': ['https://b.com']}
+    assert await kvs.get_value('CRAWLEE_STATE_0') == {'urls': ['https://a.com']}
+    assert await kvs.get_value('CRAWLEE_STATE_1') == {'urls': ['https://b.com']}
 
 
 async def test_context_handlers_use_state(key_value_store: KeyValueStore) -> None:
@@ -906,7 +909,7 @@ async def handler_three(context: BasicCrawlingContext) -> None:
     store = await crawler.get_key_value_store()
 
     # The state in the KVS must match with the last set state
-    assert (await store.get_value(f'{BasicCrawler._CRAWLEE_STATE_KEY}_0')) == {'hello': 'last_world'}
+    assert (await store.get_value('CRAWLEE_STATE_0')) == {'hello': 'last_world'}
 
 
 async def test_max_requests_per_crawl() -> None:
@@ -1334,7 +1337,7 @@ async def test_context_use_state_race_condition_in_handlers(key_value_store: Key
 
     crawler = BasicCrawler()
     store = await crawler.get_key_value_store()
-    await store.set_value(f'{BasicCrawler._CRAWLEE_STATE_KEY}_0', {'counter': 0})
+    await store.set_value('CRAWLEE_STATE_0', {'counter': 0})
     handler_barrier = Barrier(2)
 
     @crawler.router.default_handler
@@ -1349,7 +1352,7 @@ async def handler(context: BasicCrawlingContext) -> None:
     store = await crawler.get_key_value_store()
     # Ensure that local state is pushed back to kvs.
     await store.persist_autosaved_values()
-    assert (await store.get_value(f'{BasicCrawler._CRAWLEE_STATE_KEY}_0'))['counter'] == 2
+    assert (await store.get_value('CRAWLEE_STATE_0'))['counter'] == 2
 
 
 @pytest.mark.run_alone
@@ -1859,7 +1862,7 @@ async def test_crawler_state_persistence(tmp_path: Path) -> None:
         ).result()[0]
         # Expected state after first crawler run
         assert first_run_state.requests_finished == 2
-        state = await state_kvs.get_value(f'{BasicCrawler._CRAWLEE_STATE_KEY}_0')
+        state = await state_kvs.get_value('CRAWLEE_STATE_0')
         assert state.get('urls') == ['https://a.placeholder.com', 'https://b.placeholder.com']
 
     # Do not reuse the executor to simulate a fresh process to avoid modified class attributes.
@@ -1875,7 +1878,7 @@ async def test_crawler_state_persistence(tmp_path: Path) -> None:
         # 2 requests from first run and 1 request from second run.
         assert second_run_state.requests_finished == 3
 
-        state = await state_kvs.get_value(f'{BasicCrawler._CRAWLEE_STATE_KEY}_0')
+        state = await state_kvs.get_value('CRAWLEE_STATE_0')
         assert state.get('urls') == [
             'https://a.placeholder.com',
             'https://b.placeholder.com',
@@ -1912,9 +1915,9 @@ async def test_crawler_state_persistence_2_crawlers_with_migration(tmp_path: Pat
         # Expected state after first crawler run
         assert first_run_states[0].requests_finished == 1
         assert first_run_states[1].requests_finished == 1
-        state_0 = await state_kvs.get_value(f'{BasicCrawler._CRAWLEE_STATE_KEY}_0')
+        state_0 = await state_kvs.get_value('CRAWLEE_STATE_0')
         assert state_0.get('urls') == ['https://a.placeholder.com']
-        state_1 = await state_kvs.get_value(f'{BasicCrawler._CRAWLEE_STATE_KEY}_1')
+        state_1 = await state_kvs.get_value('CRAWLEE_STATE_1')
         assert state_1.get('urls') == ['https://c.placeholder.com']
 
     with ProcessPoolExecutor() as executor:
@@ -1930,9 +1933,9 @@ async def test_crawler_state_persistence_2_crawlers_with_migration(tmp_path: Pat
         # Expected state after first crawler run
         assert second_run_states[0].requests_finished == 2
         assert second_run_states[1].requests_finished == 2
-        state_0 = await state_kvs.get_value(f'{BasicCrawler._CRAWLEE_STATE_KEY}_0')
+        state_0 = await state_kvs.get_value('CRAWLEE_STATE_0')
         assert state_0.get('urls') == ['https://a.placeholder.com', 'https://b.placeholder.com']
-        state_1 = await state_kvs.get_value(f'{BasicCrawler._CRAWLEE_STATE_KEY}_1')
+        state_1 = await state_kvs.get_value('CRAWLEE_STATE_1')
         assert state_1.get('urls') == ['https://c.placeholder.com', 'https://d.placeholder.com']