Add sessions subpackage (#108)

Author: vdusek

pyproject.toml

@@ -58,6 +58,7 @@ pydantic-settings = "^2.2.1"
 pyee = "^11.1.0"
 sortedcollections = "^2.1.0"
 typing-extensions = "^4.1.0"
+python-dateutil = "^2.9.0"
 
 [tool.poetry.group.dev.dependencies]
 build = "~1.2.0"
@@ -77,6 +78,7 @@ twine = "~5.0.0"
 types-aiofiles = "^23.2.0.20240106"
 types-colorama = "~0.4.15.20240106"
 types-psutil = "~5.9.5.20240205"
+types-python-dateutil = "^2.9.0.20240316"
 
 [tool.ruff]
 line-length = 120

src/crawlee/autoscaling/snapshotter.py

@@ -94,11 +94,7 @@ def __init__(
         self._reserve_memory_ratio = reserve_memory_ratio
         self._memory_warning_cooldown_period = memory_warning_cooldown_period
         self._client_rate_limit_error_retry_count = client_rate_limit_error_retry_count
-
-        if max_memory_size is None:
-            self._max_memory_size = self._get_default_max_memory_size()
-        else:
-            self._max_memory_size = max_memory_size
+        self._max_memory_size = max_memory_size or self._get_default_max_memory_size()
 
         self._cpu_snapshots: list[CpuSnapshot] = []
         self._event_loop_snapshots: list[EventLoopSnapshot] = []

src/crawlee/events/event_manager.py

@@ -13,6 +13,7 @@
 
 if TYPE_CHECKING:
     from datetime import timedelta
+    from types import TracebackType
 
     from crawlee.events.types import Event, EventData, Listener, WrappedListener
 
@@ -26,8 +27,13 @@ class EventManager:
     their execution. It is built on top of the `pyee.asyncio.AsyncIOEventEmitter` class.
     """
 
-    def __init__(self) -> None:
-        logger.debug('Calling EventManager.__init__()...')
+    def __init__(self, close_timeout: timedelta | None = None) -> None:
+        """Create a new instance.
+
+        Args:
+            close_timeout: Optional timeout after which the pending event listeners are canceled.
+        """
+        self._close_timeout = close_timeout
 
         # Asynchronous event emitter for handle events and invoke the event listeners.
         self._event_emitter = AsyncIOEventEmitter()
@@ -41,19 +47,35 @@ def __init__(self) -> None:
             lambda: defaultdict(list),
         )
 
+    async def __aenter__(self) -> EventManager:
+        """Initializes the event manager upon entering the async context."""
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        exc_traceback: TracebackType | None,
+    ) -> None:
+        """Closes the local event manager upon exiting the async context.
+
+        This will stop listening for the events, and it will wait for all the event listeners to finish.
+        """
+        await self.wait_for_all_listeners_to_complete(timeout=self._close_timeout)
+        self._event_emitter.remove_all_listeners()
+        self._listener_tasks.clear()
+        self._listeners_to_wrappers.clear()
+
     def on(self, *, event: Event, listener: Listener) -> None:
         """Add an event listener to the event manager.
 
         Args:
             event: The Actor event for which to listen to.
             listener: The function (sync or async) which is to be called when the event is emitted.
         """
-        logger.debug('Calling EventManager.on()...')
 
         @wraps(listener)
         async def listener_wrapper(event_data: EventData) -> None:
-            logger.debug('Calling EventManager.on.listener_wrapper()...')
-
             # If the listener is a coroutine function, just call it, otherwise, run it in a separate thread
             # to avoid blocking the event loop
             coro = (
@@ -92,8 +114,6 @@ def off(self, *, event: Event, listener: Listener | None = None) -> None:
             listener: The listener which is supposed to be removed. If not passed, all listeners of this event
                 are removed.
         """
-        logger.debug('Calling EventManager.off()...')
-
         if listener:
             for listener_wrapper in self._listeners_to_wrappers[event][listener]:
                 self._event_emitter.remove_listener(event.value, listener_wrapper)
@@ -109,31 +129,15 @@ def emit(self, *, event: Event, event_data: EventData) -> None:
             event: The event which will be emitted.
             event_data: The data which will be passed to the event listeners.
         """
-        logger.debug('Calling EventManager.emit()...')
         self._event_emitter.emit(event.value, event_data)
 
-    async def close(self, *, timeout: timedelta | None = None) -> None:
-        """Close the event manager.
-
-        This will stop listening for the events, and it will wait for all the event listeners to finish.
-
-        Args:
-            timeout: Optional timeout after which the pending event listeners are canceled.
-        """
-        logger.debug('Calling EventManager.close()...')
-        await self._wait_for_all_listeners_to_complete(timeout=timeout)
-        self._event_emitter.remove_all_listeners()
-        self._listener_tasks.clear()
-        self._listeners_to_wrappers.clear()
-
-    async def _wait_for_all_listeners_to_complete(self, *, timeout: timedelta | None = None) -> None:
+    async def wait_for_all_listeners_to_complete(self, *, timeout: timedelta | None = None) -> None:
         """Wait for all currently executing event listeners to complete.
 
         Args:
             timeout: The maximum time to wait for the event listeners to finish. If they do not complete within
                 the specified timeout, they will be canceled.
         """
-        logger.debug('Calling EventManager.wait_for_all_listeners_to_complete()...')
 
         async def wait_for_listeners() -> None:
             """Gathers all listener tasks and awaits their completion, logging any exceptions encountered."""

src/crawlee/events/local_event_manager.py

@@ -25,18 +25,14 @@ def __init__(
         self,
         *,
         system_info_interval: timedelta = timedelta(seconds=60),
-        close_timeout: timedelta | None = None,
     ) -> None:
         """Create a new instance.
 
         Args:
             system_info_interval: Interval at which `SystemInfo` events are emitted.
             close_timeout: Optional timeout for closing the event manager.
         """
-        logger.debug('Calling LocalEventManager.__init__()...')
-
         self._system_info_interval = system_info_interval
-        self._close_timeout = close_timeout
 
         # Recurring task for emitting system info events.
         self._emit_system_info_event_rec_task: RecurringTask | None = None
@@ -48,12 +44,14 @@ async def __aenter__(self) -> LocalEventManager:
 
         It starts emitting system info events at regular intervals.
         """
-        logger.debug('Calling LocalEventManager.__aenter__()...')
+        await super().__aenter__()
+
         self._emit_system_info_event_rec_task = RecurringTask(
             func=self._emit_system_info_event,
             delay=self._system_info_interval,
         )
         self._emit_system_info_event_rec_task.start()
+
         return self
 
     async def __aexit__(
@@ -66,17 +64,13 @@ async def __aexit__(
 
         It stops emitting system info events and closes the event manager.
         """
-        logger.debug('Calling LocalEventManager.__aexit__()...')
-
         if self._emit_system_info_event_rec_task is not None:
             await self._emit_system_info_event_rec_task.stop()
 
-        await super().close(timeout=self._close_timeout)
+        await super().__aexit__(exc_type, exc_value, exc_traceback)
 
     async def _emit_system_info_event(self) -> None:
         """Emits a system info event with the current CPU and memory usage."""
-        logger.debug('Calling LocalEventManager._emit_system_info_event()...')
-
         cpu_info = await asyncio.to_thread(get_cpu_info)
         memory_info = await asyncio.to_thread(get_memory_info)
 

src/crawlee/events/types.py

@@ -17,6 +17,7 @@ class Event(Enum):
     MIGRATING = 'migrating'
     ABORTING = 'aborting'
     EXIT = 'exit'
+    SESSION_RETIRED = 'sessionRetired'
 
 
 @dataclass

src/crawlee/sessions/__init__.py

@@ -0,0 +1,2 @@
+from .session import Session
+from .session_pool import SessionPool

src/crawlee/sessions/models.py

@@ -0,0 +1,58 @@
+# ruff: noqa: TCH002 TCH003
+
+from __future__ import annotations
+
+from datetime import datetime, timedelta
+from typing import Annotated, Any
+
+from dateutil import parser
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+
+class SessionModel(BaseModel):
+    """Model for a Session object."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    id: Annotated[str, Field(alias='id')]
+    max_age: Annotated[timedelta, Field(alias='maxAge')]
+    user_data: Annotated[dict, Field(alias='userData')]
+    max_error_score: Annotated[float, Field(alias='maxErrorScore')]
+    error_score_decrement: Annotated[float, Field(alias='errorScoreDecrement')]
+    created_at: Annotated[datetime, Field(alias='createdAt')]
+    usage_count: Annotated[int, Field(alias='usageCount')]
+    max_usage_count: Annotated[int, Field(alias='maxUsageCount')]
+    error_score: Annotated[float, Field(alias='errorScore')]
+    cookies: Annotated[dict, Field(alias='cookies')]
+    blocked_status_codes: Annotated[list[int], Field(alias='blockedStatusCodes')]
+
+    @field_validator('max_age', mode='before')
+    @classmethod
+    def parse_max_age(cls, value: Any) -> timedelta:
+        """Try to parse max_age field into a timedelta object."""
+        if isinstance(value, timedelta):
+            return value
+
+        if isinstance(value, str):
+            try:
+                parsed_time = parser.parse(value)
+                return timedelta(hours=parsed_time.hour, minutes=parsed_time.minute, seconds=parsed_time.second)
+            except ValueError as exc:
+                raise ValueError(f"Invalid time format for max_age. Expected 'HH:MM:SS', got {value}") from exc
+
+        raise ValueError('Invalid data type for max_age')
+
+
+class SessionPoolModel(BaseModel):
+    """Model for a SessionPool object."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    persistence_enabled: Annotated[bool, Field(alias='persistenceEnabled')]
+    persist_state_kvs_name: Annotated[str, Field(alias='persistStateKvsName')]
+    persist_state_key: Annotated[str, Field(alias='persistStateKey')]
+    max_pool_size: Annotated[int, Field(alias='maxPoolSize')]
+    session_count: Annotated[int, Field(alias='sessionCount')]
+    usable_session_count: Annotated[int, Field(alias='usableSessionCount')]
+    retired_session_count: Annotated[int, Field(alias='retiredSessionCount')]
+    sessions: Annotated[list[SessionModel], Field(alias='sessions')]