Source code for lalandre_core.utils.shared_key_pool

"""Helpers for dispatching calls across a shared API key pool."""

from __future__ import annotations

from typing import Any, Callable, Dict, Mapping, TypeVar

from .api_key_pool import APIKeyPool

T = TypeVar("T")




def build_clients_by_key(
    *,
    key_pool: APIKeyPool,
    factory: Callable[[str], T],
) -> Dict[str, T]:
    """Build one client instance per key in the shared pool."""
    clients: Dict[str, T] = {}
    for key in key_pool:
        if key not in clients:
            clients[key] = factory(key)
    if not clients:
        raise ValueError("Shared key pool requires at least one client")
    return clients





class SharedKeyPoolProxy:
    """Delegate each callable access to the next client selected by ``key_pool``."""

    def __init__(
        self,
        *,
        key_pool: APIKeyPool,
        clients_by_key: Mapping[str, Any],
    ) -> None:
        if not clients_by_key:
            raise ValueError("SharedKeyPoolProxy requires at least one client")
        self._key_pool = key_pool
        self._clients_by_key = dict(clients_by_key)
        self._default_client = next(iter(self._clients_by_key.values()))

    def _select_client(self) -> Any:
        key = self._key_pool.next_key()
        client = self._clients_by_key.get(key)
        if client is None:
            raise KeyError("Shared key pool returned a key without a bound client")
        return client

    def __getattr__(self, name: str) -> Any:
        default_attr = getattr(self._default_client, name)
        if callable(default_attr):

            def _call(*args: Any, **kwargs: Any) -> Any:
                return getattr(self._select_client(), name)(*args, **kwargs)

            return _call
        return default_attr