Simplify connection config

Author: jakubno

packages/python-sdk/e2b/connection_config.py

@@ -1,7 +1,9 @@
 import os
 
 from typing import Literal, Optional, Dict
+from dataclasses import dataclass
 from httpx._types import ProxyTypes
+from typing_extensions import Unpack
 
 from e2b.api.metadata import package_version
 
@@ -11,6 +13,28 @@
 KEEPALIVE_PING_HEADER = "Keepalive-Ping-Interval"
 
 
+@dataclass
+class ApiParams:
+    """
+    Parameters for a request.
+
+    In the case of a sandbox, it applies to all **requests made to the returned sandbox**.
+    """
+
+    """Timeout for the request in **seconds**, defaults to 60 seconds."""
+    request_timeout: Optional[float]
+    """Additional headers to send with the request."""
+    headers: Dict[str, str]
+    """E2B API Key to use for authentication, defaults to `E2B_API_KEY` environment variable"""
+    api_key: Optional[str]
+    """E2B domain to use for authentication, defaults to `E2B_DOMAIN` environment variable"""
+    domain: str
+    """Whether to use debug mode, defaults to `E2B_DEBUG` environment variable"""
+    debug: bool
+    """Proxy to use for the request. In case of a sandbox it applies to all **requests made to the returned sandbox**"""
+    proxy: Optional[ProxyTypes]
+
+
 class ConnectionConfig:
     """
     Configuration for the connection to the API.
@@ -40,6 +64,7 @@ def __init__(
         access_token: Optional[str] = None,
         request_timeout: Optional[float] = None,
         headers: Optional[Dict[str, str]] = None,
+        extra_sandbox_headers: Optional[Dict[str, str]] = None,
         proxy: Optional[ProxyTypes] = None,
     ):
         self.domain = domain or ConnectionConfig._domain()
@@ -48,6 +73,7 @@ def __init__(
         self.access_token = access_token or ConnectionConfig._access_token()
         self.headers = headers or {}
         self.headers["User-Agent"] = f"e2b-python-sdk/{package_version}"
+        self.__extra_sandbox_headers = extra_sandbox_headers or {}
 
         self.proxy = proxy
 
@@ -82,6 +108,48 @@ def _get_request_timeout(
     def get_request_timeout(self, request_timeout: Optional[float] = None):
         return self._get_request_timeout(self.request_timeout, request_timeout)
 
+    def get_api_params(
+        self,
+        **opts: Unpack[ApiParams],
+    ) -> dict:
+        """
+        Get the parameters for the API call.
+
+        This is used to avoid passing the following attributes to the API call:
+        - access_token
+        - api_url
+
+        It also returns a copy, so the original object is not modified.
+
+        :return: Dictionary of parameters for the API call
+        """
+        headers = opts.get("headers")
+        request_timeout = opts.get("request_timeout")
+        api_key = opts.get("api_key")
+        domain = opts.get("domain")
+        debug = opts.get("debug")
+        proxy = opts.get("proxy")
+
+        req_headers = self.headers.copy()
+        if headers is not None:
+            req_headers.update(headers)
+
+        return ApiParams(
+            api_key=api_key or self.api_key,
+            domain=domain or self.domain,
+            debug=debug or self.debug,
+            request_timeout=self.get_request_timeout(request_timeout),
+            headers=req_headers,
+            proxy=proxy or self.proxy,
+        ).__dict__
+
+    @property
+    def sandbox_headers(self):
+        return {
+            **self.headers,
+            **self.__extra_sandbox_headers,
+        }
+
 
 Username = Literal["root", "user"]
 """

packages/python-sdk/e2b/sandbox/main.py

@@ -1,6 +1,5 @@
 import urllib.parse
 
-from abc import ABC, abstractmethod
 from typing import Optional
 
 from e2b.sandbox.signature import get_signature
@@ -9,7 +8,7 @@
 from httpx import Limits
 
 
-class SandboxSetup(ABC):
+class SandboxSetup:
     _limits = Limits(
         max_keepalive_connections=40,
         max_connections=40,
@@ -21,30 +20,48 @@ class SandboxSetup(ABC):
     default_sandbox_timeout = 300
     default_template = "base"
 
+    def __init__(
+        self,
+        sandbox_id: str,
+        envd_version: Optional[str],
+        envd_access_token: Optional[str],
+        sandbox_domain: Optional[str],
+        connection_config: ConnectionConfig,
+    ):
+        self.__connection_config = connection_config
+        self.__sandbox_id = sandbox_id
+        self.__sandbox_domain = sandbox_domain or self.connection_config.domain
+        self.__envd_version = envd_version
+        self.__envd_access_token = envd_access_token
+        self.__envd_api_url = f"{'http' if self.connection_config.debug else 'https'}://{self.get_host(self.envd_port)}"
+
+    @property
+    def _envd_access_token(self) -> Optional[str]:
+        """Private property to access the envd token"""
+        return self.__envd_access_token
+
     @property
-    @abstractmethod
     def connection_config(self) -> ConnectionConfig:
-        ...
+        return self.__connection_config
 
     @property
-    @abstractmethod
-    def _envd_access_token(self) -> Optional[str]:
-        ...
+    def _envd_version(self) -> Optional[str]:
+        return self.__envd_version
 
     @property
-    @abstractmethod
-    def envd_api_url(self) -> str:
-        ...
+    def sandbox_domain(self) -> Optional[str]:
+        return self.__sandbox_domain
 
     @property
-    @abstractmethod
-    def sandbox_id(self) -> str:
-        ...
+    def envd_api_url(self) -> str:
+        return self.__envd_api_url
 
     @property
-    @abstractmethod
-    def sandbox_domain(self) -> str:
-        ...
+    def sandbox_id(self) -> str:
+        """
+        Unique identifier of the sandbox.
+        """
+        return self.__sandbox_id
 
     def _file_url(
         self,

packages/python-sdk/e2b/sandbox/sandbox_api.py

@@ -28,7 +28,7 @@ class SandboxInfo:
     """Sandbox expiration date."""
     envd_version: Optional[str]
     """Envd version."""
-    _envd_access_token: Optional[str]
+    envd_access_token: Optional[str]
     """Envd access token."""
 
 

packages/python-sdk/e2b/sandbox_async/commands/command.py

@@ -35,7 +35,7 @@ def __init__(
             # compressor=e2b_connect.GzipCompressor,
             async_pool=pool,
             json=True,
-            headers=connection_config.headers,
+            headers=connection_config.sandbox_headers,
         )
 
     async def list(

packages/python-sdk/e2b/sandbox_async/commands/pty.py

@@ -39,6 +39,7 @@ def __init__(
             # compressor=e2b_connect.GzipCompressor,
             async_pool=pool,
             json=True,
+            headers=connection_config.sandbox_headers,
         )
 
     async def kill(

packages/python-sdk/e2b/sandbox_async/filesystem/filesystem.py

@@ -51,7 +51,7 @@ def __init__(
             # compressor=e2b_connect.GzipCompressor,
             async_pool=pool,
             json=True,
-            headers=connection_config.headers,
+            headers=connection_config.sandbox_headers,
         )
 
     @overload