Simplify connection config

Author: jakubno

packages/python-sdk/e2b/connection_config.py

@@ -1,6 +1,7 @@
 import os
 
 from typing import Literal, Optional, Dict
+from dataclasses import dataclass
 from httpx._types import ProxyTypes
 
 from e2b.api.metadata import package_version
@@ -11,6 +12,20 @@
 KEEPALIVE_PING_HEADER = "Keepalive-Ping-Interval"
 
 
+@dataclass
+class ApiParams:
+    """
+    Parameters for API calls that avoid passing access_token and api_url.
+    """
+
+    api_key: Optional[str]
+    domain: str
+    debug: bool
+    request_timeout: Optional[float]
+    headers: Dict[str, str]
+    proxy: Optional[ProxyTypes]
+
+
 class ConnectionConfig:
     """
     Configuration for the connection to the API.
@@ -40,6 +55,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 +64,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 +99,44 @@ 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,
+        request_timeout: Optional[float],
+        headers: Optional[Dict[str, str]],
+    ) -> dict:
+        """
+        Get the parameters for the API call.
+
+        This is used to avoid passing following attributes to the API call:
+        - access_token
+        - api_url
+
+        It also returns a copy so the original object is not modified.
+
+        :param request_timeout: Request timeout in seconds
+        :param headers: Additional headers to add to the request
+        :return: Dictionary of parameters for the API call
+        """
+        req_headers = self.headers.copy()
+        if headers is not None:
+            req_headers.update(headers)
+
+        return ApiParams(
+            api_key=self.api_key,
+            domain=self.domain,
+            debug=self.debug,
+            request_timeout=self.get_request_timeout(request_timeout),
+            headers=req_headers,
+            proxy=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