fix: align Python SDK migration types

Co-Authored-By: mish@e2b.dev <mish@e2b.dev>

Author: devin-ai-integration[bot]

.changeset/strong-mice-type.md

@@ -0,0 +1,5 @@
+---
+'@e2b/code-interpreter-python': patch
+---
+
+Improve Python model typing and safely handle optional chart and sandbox URL values.

python/e2b_code_interpreter/charts.py

@@ -194,17 +194,19 @@ def __init__(self, **kwargs):
 class SuperChart(Chart):
     type = ChartType.SUPERCHART
 
-    elements: List[
-        Union[LineChart, ScatterChart, BarChart, PieChart, BoxAndWhiskerChart]
-    ]
+    elements: List[Chart]
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
-        self.elements = [_deserialize_chart(g) for g in kwargs["elements"]]
+        self.elements = []
+        for raw_chart in kwargs["elements"]:
+            chart = _deserialize_chart(raw_chart)
+            if chart is not None:
+                self.elements.append(chart)
 
 
 ChartTypes = Union[
-    LineChart, ScatterChart, BarChart, PieChart, BoxAndWhiskerChart, SuperChart
+    Chart, LineChart, ScatterChart, BarChart, PieChart, BoxAndWhiskerChart, SuperChart
 ]
 
 

python/e2b_code_interpreter/code_interpreter_async.py

@@ -1,7 +1,7 @@
 import logging
 import httpx
 
-from typing import Optional, Dict, overload, Union, List
+from typing import cast, Optional, Dict, overload, Union, List
 from httpx import AsyncClient
 
 from e2b import (
@@ -63,7 +63,7 @@ class AsyncSandbox(BaseAsyncSandbox):
     def _jupyter_url(self) -> str:
         # Honors the `sandbox_url` option and the `E2B_SANDBOX_URL` environment
         # variable, same as the base SDK does for envd requests.
-        sandbox_url = self.connection_config._sandbox_url
+        sandbox_url = cast(Optional[str], self.connection_config._sandbox_url)
         if sandbox_url:
             return sandbox_url
         return f"{'http' if self.connection_config.debug else 'https'}://{self.get_host(JUPYTER_PORT)}"

python/e2b_code_interpreter/code_interpreter_sync.py

@@ -1,7 +1,7 @@
 import logging
 import httpx
 
-from typing import Optional, Dict, overload, Union, List
+from typing import cast, Optional, Dict, overload, Union, List
 from httpx import Client
 from e2b import Sandbox as BaseSandbox, InvalidArgumentException
 from e2b.api.client_sync import get_transport
@@ -60,7 +60,7 @@ class Sandbox(BaseSandbox):
     def _jupyter_url(self) -> str:
         # Honors the `sandbox_url` option and the `E2B_SANDBOX_URL` environment
         # variable, same as the base SDK does for envd requests.
-        sandbox_url = self.connection_config._sandbox_url
+        sandbox_url = cast(Optional[str], self.connection_config._sandbox_url)
         if sandbox_url:
             return sandbox_url
         return f"{'http' if self.connection_config.debug else 'https'}://{self.get_host(JUPYTER_PORT)}"

python/e2b_code_interpreter/models.py

@@ -207,7 +207,7 @@ def formats(self) -> Iterable[str]:
 
         return formats
 
-    def __str__(self) -> Optional[str]:
+    def __str__(self) -> str:
         """
         Returns the text representation of the data.
 
@@ -305,7 +305,12 @@ class Logs:
     stderr: List[str] = field(default_factory=list)
     """List of strings printed to stderr by prints, subprocesses, etc."""
 
-    def __init__(self, stdout: List[str] = None, stderr: List[str] = None, **kwargs):
+    def __init__(
+        self,
+        stdout: Optional[List[str]] = None,
+        stderr: Optional[List[str]] = None,
+        **kwargs,
+    ):
         self.stdout = stdout or []
         self.stderr = stderr or []
 
@@ -329,7 +334,9 @@ def serialize_results(results: List[Result]) -> List[Dict[str, str]]:
         serialized_dict = {}
         for key in result.formats():
             if key == "chart":
-                serialized_dict[key] = result.chart.to_dict()
+                chart = result.chart
+                if chart is not None:
+                    serialized_dict[key] = chart.to_dict()
             else:
                 serialized_dict[key] = result[key]
 
@@ -356,8 +363,8 @@ class Execution:
 
     def __init__(
         self,
-        results: List[Result] = None,
-        logs: Logs = None,
+        results: Optional[List[Result]] = None,
+        logs: Optional[Logs] = None,
         error: Optional[ExecutionError] = None,
         execution_count: Optional[int] = None,
         **kwargs,
@@ -510,7 +517,7 @@ def __init__(self, context_id: str, language: str, cwd: str, **kwargs):
     @classmethod
     def from_json(cls, data: Dict[str, str]):
         return cls(
-            context_id=data.get("id"),
-            language=data.get("language"),
-            cwd=data.get("cwd"),
+            context_id=data["id"],
+            language=data["language"],
+            cwd=data["cwd"],
         )