Fix Figure import bug

e2b-dev · Aug 28, 2024 · 14dce4c · 14dce4c
1 parent 6e041b9
commit 14dce4c
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 7 deletions.
diff --git a/python/tests/async/test_async_env_vars.py b/python/tests/async/test_async_env_vars.py
@@ -9,14 +9,20 @@ async def test_env_vars_sandbox():
 
 
 async def test_env_vars_in_exec_cell(async_sandbox: AsyncCodeInterpreter):
-    result = await async_sandbox.notebook.exec_cell("import os; os.getenv('FOO')", envs={"FOO": "bar"})
+    result = await async_sandbox.notebook.exec_cell(
+        "import os; os.getenv('FOO')", envs={"FOO": "bar"}
+    )
     assert result.text == "bar"
 
 
 async def test_env_vars_override():
     sbx = await AsyncCodeInterpreter.create(envs={"FOO": "bar", "SBX": "value"})
-    await sbx.notebook.exec_cell("import os; os.environ['FOO'] = 'bar'; os.environ['RUNTIME_ENV'] = 'value'")
-    result = await sbx.notebook.exec_cell("import os; os.getenv('FOO')", envs={"FOO": "baz"})
+    await sbx.notebook.exec_cell(
+        "import os; os.environ['FOO'] = 'bar'; os.environ['RUNTIME_ENV'] = 'value'"
+    )
+    result = await sbx.notebook.exec_cell(
+        "import os; os.getenv('FOO')", envs={"FOO": "baz"}
+    )
     assert result.text == "baz"
 
     result = await sbx.notebook.exec_cell("import os; os.getenv('RUNTIME_ENV')")
@@ -29,4 +35,3 @@ async def test_env_vars_override():
     assert result.text == "bar"
 
     await sbx.kill()
-
diff --git a/python/tests/sync/test_env_vars.py b/python/tests/sync/test_env_vars.py
@@ -9,13 +9,17 @@ async def test_env_vars_sandbox():
 
 
 async def test_env_vars_in_exec_cell(sandbox: CodeInterpreter):
-    result = sandbox.notebook.exec_cell("import os; os.getenv('FOO')", envs={"FOO": "bar"})
+    result = sandbox.notebook.exec_cell(
+        "import os; os.getenv('FOO')", envs={"FOO": "bar"}
+    )
     assert result.text == "bar"
 
 
 async def test_env_vars_override():
     sbx = CodeInterpreter(envs={"FOO": "bar", "SBX": "value"})
-    sbx.notebook.exec_cell("import os; os.environ['FOO'] = 'bar'; os.environ['RUNTIME_ENV'] = 'value'")
+    sbx.notebook.exec_cell(
+        "import os; os.environ['FOO'] = 'bar'; os.environ['RUNTIME_ENV'] = 'value'"
+    )
     result = sbx.notebook.exec_cell("import os; os.getenv('FOO')", envs={"FOO": "baz"})
     assert result.text == "baz"
 

diff --git a/template/ipython_startup_script.py b/template/ipython_startup_script.py
@@ -106,8 +106,11 @@ class E2BFormatter(BaseFormatter):
     type_printers = {pandas.DataFrame: _data_frame_repr_e2b_data_}
 
     def __call__(self, obj):
+        # Figure object is for some reason removed on execution of the cell,
+        # so it can't be used in type_printers or with top-level import
+        from matplotlib.pyplot import Figure
+
         if isinstance(obj, Figure):
-            # Figure object is for some reason removed from type_printers
             return _figure_repr_e2b_data_(obj)
         return super().__call__(obj)