feat: Add use_state context method (#682)

### Description

Add `use_state` context helper method with auto-save behavior inspired
by TS's realization of `crawlee`.
Adds a `get_auto_saved_value` method for `kvs`. Working with internal
cache and saving data on `PERSIST_STATE` event

(At this point, it's more of a gut check and synchronization with the
team. It will probably require some more refinement)

### Issues

- Closes: #191 

### Testing

For correct testing, it is necessary to use event_manager in tests for
KeyValueStore, therefore I have added a new fixture.

### Checklist

- [x] CI passed

---------

Co-authored-by: Vlada Dusek <[email protected]>

Author: Mantisus

src/crawlee/_types.py

@@ -300,6 +300,20 @@ def __call__(
     ) -> Coroutine[None, None, HttpResponse]: ...
 
 
+class UseStateFunction(Protocol):
+    """Type of a function for performing use state.
+
+    Warning:
+        This is an experimental feature. The behavior and interface may change in future versions.
+    """
+
+    def __call__(
+        self,
+        key: str,
+        default_value: dict[str, JsonSerializable] | None = None,
+    ) -> Coroutine[None, None, dict[str, JsonSerializable]]: ...
+
+
 T = TypeVar('T')
 
 
@@ -347,6 +361,7 @@ class BasicCrawlingContext:
     send_request: SendRequestFunction
     add_requests: AddRequestsFunction
     push_data: PushDataFunction
+    use_state: UseStateFunction
     get_key_value_store: GetKeyValueStoreFromRequestHandlerFunction
     log: logging.Logger
 

src/crawlee/basic_crawler/_basic_crawler.py

@@ -478,6 +478,8 @@ def sigint_handler() -> None:
         self._running = False
         self._has_finished_before = True
 
+        await self._save_crawler_state()
+
         final_statistics = self._statistics.calculate()
         self._logger.info(f'Final request statistics:\n{final_statistics.to_table()}')
 
@@ -532,6 +534,16 @@ 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, key: str, default_value: dict[str, JsonSerializable] | None = None
+    ) -> dict[str, JsonSerializable]:
+        store = await self.get_key_value_store()
+        return await store.get_auto_saved_value(key, default_value)
+
+    async def _save_crawler_state(self) -> None:
+        store = await self.get_key_value_store()
+        await store.persist_autosaved_values()
+
     async def get_data(
         self,
         dataset_id: str | None = None,
@@ -953,6 +965,7 @@ async def __run_task_function(self) -> None:
             add_requests=result.add_requests,
             push_data=result.push_data,
             get_key_value_store=result.get_key_value_store,
+            use_state=self._use_state,
             log=self._logger,
         )
 

src/crawlee/playwright_crawler/_playwright_crawler.py

@@ -136,6 +136,7 @@ async def _open_page(self, context: BasicCrawlingContext) -> AsyncGenerator[Play
             add_requests=context.add_requests,
             send_request=context.send_request,
             push_data=context.push_data,
+            use_state=context.use_state,
             proxy_info=context.proxy_info,
             get_key_value_store=context.get_key_value_store,
             log=context.log,
@@ -225,6 +226,7 @@ async def enqueue_links(
                 add_requests=context.add_requests,
                 send_request=context.send_request,
                 push_data=context.push_data,
+                use_state=context.use_state,
                 proxy_info=context.proxy_info,
                 get_key_value_store=context.get_key_value_store,
                 log=context.log,

src/crawlee/storages/_key_value_store.py

@@ -1,16 +1,19 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, TypeVar, overload
+from typing import TYPE_CHECKING, Any, ClassVar, TypeVar, overload
 
 from typing_extensions import override
 
+from crawlee import service_container
 from crawlee._utils.docs import docs_group
 from crawlee.base_storage_client._models import KeyValueStoreKeyInfo, KeyValueStoreMetadata
+from crawlee.events._types import Event, EventPersistStateData
 from crawlee.storages._base_storage import BaseStorage
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator
 
+    from crawlee._types import JsonSerializable
     from crawlee.base_storage_client import BaseStorageClient
     from crawlee.configuration import Configuration
 
@@ -52,6 +55,10 @@ class KeyValueStore(BaseStorage):
     ```
     """
 
+    # Cache for persistent (auto-saved) values
+    _general_cache: ClassVar[dict[str, dict[str, dict[str, JsonSerializable]]]] = {}
+    _persist_state_event_started = False
+
     def __init__(
         self,
         id: str,
@@ -105,6 +112,7 @@ async def drop(self) -> None:
         from crawlee.storages._creation_management import remove_storage_from_cache
 
         await self._resource_client.delete()
+        self._clear_cache()
         remove_storage_from_cache(storage_class=self.__class__, id=self._id, name=self._name)
 
     @overload
@@ -175,3 +183,70 @@ async def get_public_url(self, key: str) -> str:
             The public URL for the given key.
         """
         return await self._resource_client.get_public_url(key)
+
+    async def get_auto_saved_value(
+        self, key: str, default_value: dict[str, JsonSerializable] | None = None
+    ) -> dict[str, JsonSerializable]:
+        """Gets a value from KVS that will be automatically saved on changes.
+
+        Args:
+            key: Key of the record, to store the value.
+            default_value: Value to be used if the record does not exist yet. Should be a dictionary.
+
+        Returns:
+            Returns the value of the key.
+        """
+        default_value = {} if default_value is None else default_value
+
+        if key in self._cache:
+            return self._cache[key]
+
+        value = await self.get_value(key, default_value)
+
+        if not isinstance(value, dict):
+            raise TypeError(
+                f'Expected dictionary for persist state value at key "{key}, but got {type(value).__name__}'
+            )
+
+        self._cache[key] = value
+
+        self._ensure_persist_event()
+
+        return value
+
+    @property
+    def _cache(self) -> dict[str, dict[str, JsonSerializable]]:
+        """Cache dictionary for storing auto-saved values indexed by store ID."""
+        if self._id not in self._general_cache:
+            self._general_cache[self._id] = {}
+        return self._general_cache[self._id]
+
+    async def _persist_save(self, _event_data: EventPersistStateData | None = None) -> None:
+        """Save cache with persistent values. Can be used in Event Manager."""
+        for key, value in self._cache.items():
+            await self.set_value(key, value)
+
+    def _ensure_persist_event(self) -> None:
+        """Setup persist state event handling if not already done."""
+        if self._persist_state_event_started:
+            return
+
+        event_manager = service_container.get_event_manager()
+        event_manager.on(event=Event.PERSIST_STATE, listener=self._persist_save)
+        self._persist_state_event_started = True
+
+    def _clear_cache(self) -> None:
+        """Clear cache with persistent values."""
+        self._cache.clear()
+
+    def _drop_persist_state_event(self) -> None:
+        """Off event_manager listener and drop event status."""
+        if self._persist_state_event_started:
+            event_manager = service_container.get_event_manager()
+            event_manager.off(event=Event.PERSIST_STATE, listener=self._persist_save)
+        self._persist_state_event_started = False
+
+    async def persist_autosaved_values(self) -> None:
+        """Force persistent values to be saved without waiting for an event in Event Manager."""
+        if self._persist_state_event_started:
+            await self._persist_save()

tests/unit/basic_crawler/test_basic_crawler.py

@@ -24,11 +24,20 @@
 from crawlee.storages import Dataset, KeyValueStore, RequestList, RequestQueue
 
 if TYPE_CHECKING:
-    from collections.abc import Sequence
+    from collections.abc import AsyncGenerator, Sequence
 
     import respx
     from yarl import URL
 
+    from crawlee._types import JsonSerializable
+
+
+@pytest.fixture
+async def key_value_store() -> AsyncGenerator[KeyValueStore, None]:
+    kvs = await KeyValueStore.open()
+    yield kvs
+    await kvs.drop()
+
 
 async def test_processes_requests() -> None:
     crawler = BasicCrawler(request_provider=RequestList(['http://a.com/', 'http://b.com/', 'http://c.com/']))
@@ -677,6 +686,63 @@ async def handler(context: BasicCrawlingContext) -> None:
     assert (await store.get_value('foo')) == 'bar'
 
 
+async def test_context_use_state(key_value_store: KeyValueStore) -> None:
+    crawler = BasicCrawler()
+
+    @crawler.router.default_handler
+    async def handler(context: BasicCrawlingContext) -> None:
+        await context.use_state('state', {'hello': 'world'})
+
+    await crawler.run(['https://hello.world'])
+
+    store = await crawler.get_key_value_store()
+
+    assert (await store.get_value('state')) == {'hello': 'world'}
+
+
+async def test_context_handlers_use_state(key_value_store: KeyValueStore) -> None:
+    state_in_handler_one: dict[str, JsonSerializable] = {}
+    state_in_handler_two: dict[str, JsonSerializable] = {}
+    state_in_handler_three: dict[str, JsonSerializable] = {}
+
+    crawler = BasicCrawler()
+
+    @crawler.router.handler('one')
+    async def handler_one(context: BasicCrawlingContext) -> None:
+        state = await context.use_state('state', {'hello': 'world'})
+        state_in_handler_one.update(state)
+        state['hello'] = 'new_world'
+        await context.add_requests([Request.from_url('https://crawlee.dev/docs/quick-start', label='two')])
+
+    @crawler.router.handler('two')
+    async def handler_two(context: BasicCrawlingContext) -> None:
+        state = await context.use_state('state', {'hello': 'world'})
+        state_in_handler_two.update(state)
+        state['hello'] = 'last_world'
+
+    @crawler.router.handler('three')
+    async def handler_three(context: BasicCrawlingContext) -> None:
+        state = await context.use_state('state', {'hello': 'world'})
+        state_in_handler_three.update(state)
+
+    await crawler.run([Request.from_url('https://crawlee.dev/', label='one')])
+    await crawler.run([Request.from_url('https://crawlee.dev/docs/examples', label='three')])
+
+    # The state in handler_one must match the default state
+    assert state_in_handler_one == {'hello': 'world'}
+
+    # The state in handler_two must match the state updated in handler_one
+    assert state_in_handler_two == {'hello': 'new_world'}
+
+    # The state in handler_three must match the final state updated in previous run
+    assert state_in_handler_three == {'hello': 'last_world'}
+
+    store = await crawler.get_key_value_store()
+
+    # The state in the KVS must match with the last set state
+    assert (await store.get_value('state')) == {'hello': 'last_world'}
+
+
 async def test_max_requests_per_crawl(httpbin: URL) -> None:
     start_urls = [
         str(httpbin / '1'),

tests/unit/basic_crawler/test_context_pipeline.py

@@ -38,6 +38,7 @@ async def test_calls_consumer_without_middleware() -> None:
         session=Session(),
         proxy_info=AsyncMock(),
         push_data=AsyncMock(),
+        use_state=AsyncMock(),
         get_key_value_store=AsyncMock(),
         log=logging.getLogger(),
     )
@@ -64,6 +65,7 @@ async def middleware_a(context: BasicCrawlingContext) -> AsyncGenerator[Enhanced
             session=context.session,
             proxy_info=AsyncMock(),
             push_data=AsyncMock(),
+            use_state=AsyncMock(),
             get_key_value_store=AsyncMock(),
             log=logging.getLogger(),
         )
@@ -80,6 +82,7 @@ async def middleware_b(context: EnhancedCrawlingContext) -> AsyncGenerator[MoreE
             session=context.session,
             proxy_info=AsyncMock(),
             push_data=AsyncMock(),
+            use_state=AsyncMock(),
             get_key_value_store=AsyncMock(),
             log=logging.getLogger(),
         )
@@ -94,6 +97,7 @@ async def middleware_b(context: EnhancedCrawlingContext) -> AsyncGenerator[MoreE
         session=Session(),
         proxy_info=AsyncMock(),
         push_data=AsyncMock(),
+        use_state=AsyncMock(),
         get_key_value_store=AsyncMock(),
         log=logging.getLogger(),
     )
@@ -119,6 +123,7 @@ async def test_wraps_consumer_errors() -> None:
         session=Session(),
         proxy_info=AsyncMock(),
         push_data=AsyncMock(),
+        use_state=AsyncMock(),
         get_key_value_store=AsyncMock(),
         log=logging.getLogger(),
     )
@@ -147,6 +152,7 @@ async def step_2(context: BasicCrawlingContext) -> AsyncGenerator[BasicCrawlingC
         session=Session(),
         proxy_info=AsyncMock(),
         push_data=AsyncMock(),
+        use_state=AsyncMock(),
         get_key_value_store=AsyncMock(),
         log=logging.getLogger(),
     )
@@ -178,6 +184,7 @@ async def step_2(context: BasicCrawlingContext) -> AsyncGenerator[BasicCrawlingC
         session=Session(),
         proxy_info=AsyncMock(),
         push_data=AsyncMock(),
+        use_state=AsyncMock(),
         get_key_value_store=AsyncMock(),
         log=logging.getLogger(),
     )