Change to TypedDict

Author: Pijukatel

src/apify/_configuration.py

@@ -7,8 +7,8 @@
 from pathlib import Path
 from typing import Annotated, Any
 
-from pydantic import AliasChoices, BaseModel, BeforeValidator, Field, model_validator
-from typing_extensions import Self, deprecated
+from pydantic import AliasChoices, BeforeValidator, Field, model_validator
+from typing_extensions import Self, TypedDict, deprecated
 
 from crawlee import service_locator
 from crawlee._utils.models import timedelta_ms
@@ -34,26 +34,26 @@ def _transform_to_list(value: Any) -> list[str] | None:
     return value if isinstance(value, list) else str(value).split(',')
 
 
-class ActorStorages(BaseModel):
-    """Storage IDs for different storage types used by an Actor."""
-
+class ActorStorages(TypedDict):
     key_value_stores: dict[str, str]
     datasets: dict[str, str]
     request_queues: dict[str, str]
 
 
-def _load_storage_keys(data: None | str | dict | ActorStorages) -> ActorStorages | None:
+def _load_storage_keys(
+    data: None | str | ActorStorages,
+) -> ActorStorages | None:
     """Load storage keys from environment."""
     if data is None:
         return None
-    if isinstance(data, ActorStorages):
-        return data
-    storage_mapping = data if isinstance(data, dict) else json.loads(data)
-    return ActorStorages(
-        key_value_stores=storage_mapping.get('keyValueStores', {}),
-        datasets=storage_mapping.get('datasets', {}),
-        request_queues=storage_mapping.get('requestQueues', {}),
-    )
+    if isinstance(data, str):
+        storage_mapping = json.loads(data)
+        return {
+            'key_value_stores': storage_mapping.get('keyValueStores', {}),
+            'datasets': storage_mapping.get('datasets', {}),
+            'request_queues': storage_mapping.get('requestQueues', {}),
+        }
+    return data
 
 
 @docs_group('Configuration')

src/apify/storage_clients/_apify/_alias_resolving.py

@@ -197,9 +197,9 @@ async def resolve_id(self) -> str | None:
         # This mapping is maintained by the Apify platform and does not have to be maintained in the default KVS.
         if self._configuration.actor_storages and self._alias != 'default':
             storage_maps = {
-                'Dataset': self._configuration.actor_storages.datasets,
-                'KeyValueStore': self._configuration.actor_storages.key_value_stores,
-                'RequestQueue': self._configuration.actor_storages.request_queues,
+                'Dataset': self._configuration.actor_storages['datasets'],
+                'KeyValueStore': self._configuration.actor_storages['key_value_stores'],
+                'RequestQueue': self._configuration.actor_storages['request_queues'],
             }
             if storage_id := storage_maps.get(self._storage_type, {}).get(self._alias):
                 return storage_id

tests/e2e/test_schema_storages/actor_source/main.py

@@ -4,4 +4,4 @@
 async def main() -> None:
     async with Actor:
         assert Actor.configuration.actor_storages
-        assert (await Actor.open_dataset(alias='custom')).id == Actor.configuration.actor_storages.datasets['custom']
+        assert (await Actor.open_dataset(alias='custom')).id == Actor.configuration.actor_storages['datasets']['custom']

tests/unit/actor/test_configuration.py

@@ -318,6 +318,6 @@ def test_actor_storage_json_env_var(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setenv('ACTOR_STORAGES_JSON', actor_storages_json)
     config = ApifyConfiguration()
     assert config.actor_storages
-    assert config.actor_storages.datasets == datasets
-    assert config.actor_storages.request_queues == request_queues
-    assert config.actor_storages.key_value_stores == key_value_stores
+    assert config.actor_storages['datasets'] == datasets
+    assert config.actor_storages['request_queues'] == request_queues
+    assert config.actor_storages['key_value_stores'] == key_value_stores

tests/unit/storage_clients/test_alias_resolver.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from apify._configuration import ActorStorages, Configuration
+from apify._configuration import Configuration
 from apify.storage_clients._apify._alias_resolving import AliasResolver
 
 
@@ -88,9 +88,7 @@ async def test_configuration_storages_alias_resolving() -> None:
 
     # Set up the configuration with the storage mapping
     configuration = Configuration(
-        actor_storages=ActorStorages(
-            datasets=datasets, request_queues=request_queues, key_value_stores=key_value_stores
-        ),
+        actor_storages={'datasets': datasets, 'request_queues': request_queues, 'key_value_stores': key_value_stores}
     )
 
     # Check that id of each non-default storage saved in the mapping is resolved