Return images

Author: jakubno

js/tests/images/bar.test.ts

@@ -0,0 +1,55 @@
+import { expect } from 'vitest'
+
+import { sandboxTest } from '../setup'
+
+// Skip this test if we are running in debug mode — the pwd and user in the testing docker container are not the same as in the actual sandbox.
+sandboxTest('test show image', async ({ sandbox }) => {
+  const code = `
+    import numpy
+    from PIL import Image
+    
+    imarray = numpy.random.rand(16,16,3) * 255
+    im = Image.fromarray(imarray.astype('uint8')).convert('RGBA')
+    
+    image.show()
+    `
+
+  const execution = await sandbox.runCode(code)
+
+  const image = execution.results[0].png
+  expect(image).toBeDefined()
+})
+
+sandboxTest('test image represent', async ({ sandbox }) => {
+  const code = `
+    import numpy
+    from PIL import Image
+
+    imarray = numpy.random.rand(16,16,3) * 255
+    im = Image.fromarray(imarray.astype('uint8')).convert('RGBA')
+
+    image
+    `
+  const execution = await sandbox.runCode(code)
+
+  const image = execution.results[0].png
+  expect(image).toBeDefined()
+})
+
+
+ sandboxTest('get image on save', async ({ sandbox }) => {
+   const code = `
+    import numpy
+    from PIL import Image
+
+    imarray = numpy.random.rand(16,16,3) * 255
+    im = Image.fromarray(imarray.astype('uint8')).convert('RGBA')
+
+    image.save("test.png")
+    `
+
+   const execution = await sandbox.runCode(code)
+
+   const image = execution.results[0].png
+   expect(image).toBeDefined()
+ })

python/tests/images/test_images.py

@@ -0,0 +1,51 @@
+from e2b_code_interpreter.code_interpreter_async import AsyncSandbox
+
+
+async def test_show_image(async_sandbox: AsyncSandbox):
+    code = """
+    import numpy
+    from PIL import Image
+    
+    imarray = numpy.random.rand(16,16,3) * 255
+    im = Image.fromarray(imarray.astype('uint8')).convert('RGBA')
+    
+    image.show()
+    """
+
+    execution = await async_sandbox.run_code(code)
+
+    image = execution.results[0].png
+    assert image
+
+
+async def test_image_as_last_command(async_sandbox: AsyncSandbox):
+    code = """
+    import numpy
+    from PIL import Image
+    
+    imarray = numpy.random.rand(16,16,3) * 255
+    im = Image.fromarray(imarray.astype('uint8')).convert('RGBA')
+    
+    image
+    """
+    execution = await async_sandbox.run_code(code)
+
+    image = execution.results[0].png
+    assert image
+
+
+async def test_get_image_on_save(async_sandbox: AsyncSandbox):
+    code = """
+    import numpy
+    from PIL import Image
+    
+    imarray = numpy.random.rand(16,16,3) * 255
+    im = Image.fromarray(imarray.astype('uint8')).convert('RGBA')
+    
+    image.save("test.png")
+    """
+
+    execution = await async_sandbox.run_code(code)
+
+    image = execution.results[0].png
+    assert image

template/startup_scripts/0003_images.py

@@ -0,0 +1,43 @@
+from typing import Any, io
+
+from IPython.core.display_functions import display
+from PIL import ImageShow
+from PIL.Image import Image
+
+original_show = ImageShow.show
+
+
+def show_file(self, path: str, **options: Any) -> int:
+    # To prevent errors from trying to display image without any display
+    return 0
+
+
+ImageShow.show_file = show_file
+
+original_save = Image.save
+
+
+# To prevent circular save and display calls
+def __repr_image(self, image_format: str, **kwargs: Any) -> bytes | None:
+    """Helper function for iPython display hook.
+
+    :param image_format: Image format.
+    :returns: image as bytes, saved into the given format.
+    """
+    b = io.BytesIO()
+    try:
+        original_save(self, b, image_format, **kwargs)
+    except Exception:
+        return None
+    return b.getvalue()
+
+
+Image._repr_image = __repr_image
+
+
+def save(image, fp, format=None, **options):
+    display(image)
+    original_save(image, fp, format, **options)
+
+
+Image.save = save