__init__.py

import asyncio
import weakref
from typing import Dict, Optional, Tuple

from httpx._types import ProxyTypes

from e2b.api import (
    _TCPKeepaliveAsyncHTTPTransport,
    AsyncApiClient,
    connection_retries,
    limits,
)
from e2b.connection_config import ConnectionConfig

TransportKey = Tuple[bool, Optional[ProxyTypes]]


def get_api_client(config: ConnectionConfig, **kwargs) -> AsyncApiClient:
    return AsyncApiClient(
        config,
        async_transport_factory=lambda: get_transport(config),
        **kwargs,
    )


class AsyncTransportWithLogger(_TCPKeepaliveAsyncHTTPTransport):
    # Keyed weakly by the event loop object itself, not id(loop) — CPython
    # reuses object ids, so a new loop could otherwise inherit a transport
    # bound to a previous, closed loop.
    _instances: weakref.WeakKeyDictionary[
        asyncio.AbstractEventLoop,
        Dict[TransportKey, "AsyncTransportWithLogger"],
    ] = weakref.WeakKeyDictionary()

    @property
    def pool(self):
        return self._pool


def _create_transport(cls, config: ConnectionConfig, http2: bool):
    """Build a keepalive transport of the given class for this config."""
    return cls(
        limits=limits,
        proxy=config.proxy,
        http2=http2,
        retries=connection_retries,
    )


def _get_cached_transport(cls, config: ConnectionConfig, http2: bool):
    loop = asyncio.get_running_loop()
    loop_instances = cls._instances.get(loop)
    if loop_instances is None:
        loop_instances = {}
        cls._instances[loop] = loop_instances

    key: TransportKey = (http2, config.proxy)
    transport = loop_instances.get(key)
    if transport is None:
        transport = _create_transport(cls, config, http2)
        loop_instances[key] = transport

    return transport


def get_transport(
    config: ConnectionConfig, http2: bool = True
) -> AsyncTransportWithLogger:
    return _get_cached_transport(AsyncTransportWithLogger, config, http2)


class AsyncEnvdTransportWithLogger(AsyncTransportWithLogger):
    _instances: weakref.WeakKeyDictionary[
        asyncio.AbstractEventLoop,
        Dict[TransportKey, "AsyncEnvdTransportWithLogger"],
    ] = weakref.WeakKeyDictionary()


def get_envd_transport(
    config: ConnectionConfig, http2: bool = True
) -> AsyncEnvdTransportWithLogger:
    return _get_cached_transport(AsyncEnvdTransportWithLogger, config, http2)