Move to models

Author: Pijukatel

src/apify/storage_clients/_apify/_models.py

@@ -5,7 +5,7 @@
 
 from pydantic import BaseModel, ConfigDict, Field
 
-from crawlee.storage_clients.models import KeyValueStoreMetadata
+from crawlee.storage_clients.models import KeyValueStoreMetadata, RequestQueueMetadata
 
 from apify import Request
 from apify._utils import docs_group
@@ -105,3 +105,27 @@ class CachedRequest(BaseModel):
 
     lock_expires_at: datetime | None = None
     """The expiration time of the lock on the request."""
+
+
+class RequestQueueStats(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+
+    delete_count: Annotated[int, Field(alias='deleteCount', default=0)]
+    """"The number of request queue deletes."""
+
+    head_item_read_count: Annotated[int, Field(alias='headItemReadCount', default=0)]
+    """The number of request queue head reads."""
+
+    read_count: Annotated[int, Field(alias='readCount', default=0)]
+    """The number of request queue reads."""
+
+    storage_bytes: Annotated[int, Field(alias='storageBytes', default=0)]
+    """Storage size in Bytes."""
+
+    write_count: Annotated[int, Field(alias='writeCount', default=0)]
+    """The number of request queue writes."""
+
+
+class ApifyRequestQueueMetadata(RequestQueueMetadata):
+    stats: Annotated[RequestQueueStats, Field(alias='stats', default_factory=RequestQueueStats)]
+    """Additional statistics about the request queue."""

src/apify/storage_clients/_apify/_request_queue_client.py

@@ -7,18 +7,23 @@
 from datetime import datetime, timedelta, timezone
 from hashlib import sha256
 from logging import getLogger
-from typing import TYPE_CHECKING, Annotated, Final
+from typing import TYPE_CHECKING, Final
 
 from cachetools import LRUCache
-from pydantic import BaseModel, ConfigDict, Field
 from typing_extensions import override
 
 from apify_client import ApifyClientAsync
 from crawlee._utils.crypto import crypto_random_object_id
 from crawlee.storage_clients._base import RequestQueueClient
 from crawlee.storage_clients.models import AddRequestsResponse, ProcessedRequest, RequestQueueMetadata
 
-from ._models import CachedRequest, ProlongRequestLockResponse, RequestQueueHead
+from ._models import (
+    ApifyRequestQueueMetadata,
+    CachedRequest,
+    ProlongRequestLockResponse,
+    RequestQueueHead,
+    RequestQueueStats,
+)
 from apify import Request
 
 if TYPE_CHECKING:
@@ -54,30 +59,6 @@ def unique_key_to_request_id(unique_key: str, *, request_id_length: int = 15) ->
     return url_safe_key[:request_id_length]
 
 
-class RequestQueueStats(BaseModel):
-    model_config = ConfigDict(populate_by_name=True)
-
-    delete_count: Annotated[int, Field(alias='deleteCount', default=0)]
-    """"The number of request queue deletes."""
-
-    head_item_read_count: Annotated[int, Field(alias='headItemReadCount', default=0)]
-    """The number of request queue head reads."""
-
-    read_count: Annotated[int, Field(alias='readCount', default=0)]
-    """The number of request queue reads."""
-
-    storage_bytes: Annotated[int, Field(alias='storageBytes', default=0)]
-    """Storage size in Bytes."""
-
-    write_count: Annotated[int, Field(alias='writeCount', default=0)]
-    """The number of request queue writes."""
-
-
-class ApifyRequestQueueMetadata(RequestQueueMetadata):
-    stats: Annotated[RequestQueueStats, Field(alias='stats', default_factory=RequestQueueStats)]
-    """Additional statistics about the request queue."""
-
-
 class ApifyRequestQueueClient(RequestQueueClient):
     """An Apify platform implementation of the request queue client."""
 

tests/integration/test_request_queue.py

@@ -14,7 +14,7 @@
     from crawlee.storages import RequestQueue
 
     from .conftest import MakeActorFunction, RunActorFunction
-    from apify.storage_clients._apify._request_queue_client import ApifyRequestQueueMetadata
+    from apify.storage_clients._apify._models import ApifyRequestQueueMetadata
 
 
 async def test_add_and_fetch_requests(