aiodataloader

import asyncio from aiodataloader import DataLoader # A mock batch loading function for demonstration async def fetch_users_from_db(user_ids: list[int]) -> list[dict | None]: print(f"Fetching users with IDs: {user_ids}") # Simulate an async database call await asyncio.sleep(0.01) # In a real scenario, this would query a database (e.g., ORM, API) users_data = { 1: {"id": 1, "name": "Alice"}, 2: {"id": 2, "name": "Bob"}, 3: {"id": 3, "name": "Charlie"}, } # Important: return values in the same order as keys, with None for missing return [users_data.get(uid) for uid in user_ids] class UserLoader(DataLoader): def __init__(self): super().__init__(self.batch_load_fn) async def batch_load_fn(self, keys: list[int]) -> list[dict | None]: return await fetch_users_from_db(keys) async def main(): user_loader = UserLoader() # Load individual users concurrently # These three loads will be coalesced into a single call to fetch_users_from_db user1_task = user_loader.load(1) user2_task = user_loader.load(2) user3_task = user_loader.load(1) # This will be served from cache (for ID 1) from the first load user1, user2, user3 = await asyncio.gather(user1_task, user2_task, user3_task) print(f"User 1 (from first load): {user1}") print(f"User 2: {user2}") print(f"User 3 (from cache): {user3}") # Example of loading many users_many = await user_loader.load_many([2, 3, 4]) # ID 4 will result in None print(f"Users (many): {users_many}") if __name__ == "__main__": asyncio.run(main())