Fix image generation

Author: jakubno

js/tests/images/bar.test.ts

@@ -9,9 +9,10 @@ sandboxTest('test show image', async ({ sandbox }) => {
     from PIL import Image
     
     imarray = numpy.random.rand(16,16,3) * 255
-    im = Image.fromarray(imarray.astype('uint8')).convert('RGBA')
+    image = Image.fromarray(imarray.astype('uint8')).convert('RGBA')
     
     image.show()
+    print("done")
     `
 
   const execution = await sandbox.runCode(code)
@@ -26,7 +27,7 @@ sandboxTest('test image represent', async ({ sandbox }) => {
     from PIL import Image
 
     imarray = numpy.random.rand(16,16,3) * 255
-    im = Image.fromarray(imarray.astype('uint8')).convert('RGBA')
+    image = Image.fromarray(imarray.astype('uint8')).convert('RGBA')
 
     image
     `
@@ -43,7 +44,7 @@ sandboxTest('test image represent', async ({ sandbox }) => {
     from PIL import Image
 
     imarray = numpy.random.rand(16,16,3) * 255
-    im = Image.fromarray(imarray.astype('uint8')).convert('RGBA')
+    image = Image.fromarray(imarray.astype('uint8')).convert('RGBA')
 
     image.save("test.png")
     `

python/tests/images/test_images.py

@@ -7,9 +7,10 @@ async def test_show_image(async_sandbox: AsyncSandbox):
     from PIL import Image
     
     imarray = numpy.random.rand(16,16,3) * 255
-    im = Image.fromarray(imarray.astype('uint8')).convert('RGBA')
+    image = Image.fromarray(imarray.astype('uint8')).convert('RGBA')
     
     image.show()
+    print("Image shown.")
     """
 
     execution = await async_sandbox.run_code(code)
@@ -24,7 +25,7 @@ async def test_image_as_last_command(async_sandbox: AsyncSandbox):
     from PIL import Image
     
     imarray = numpy.random.rand(16,16,3) * 255
-    im = Image.fromarray(imarray.astype('uint8')).convert('RGBA')
+    image = Image.fromarray(imarray.astype('uint8')).convert('RGBA')
     
     image
     """
@@ -40,9 +41,10 @@ async def test_get_image_on_save(async_sandbox: AsyncSandbox):
     from PIL import Image
     
     imarray = numpy.random.rand(16,16,3) * 255
-    im = Image.fromarray(imarray.astype('uint8')).convert('RGBA')
+    image = Image.fromarray(imarray.astype('uint8')).convert('RGBA')
     
     image.save("test.png")
+    print("Image saved.")
     """
 
     execution = await async_sandbox.run_code(code)

template/startup_scripts/0003_images.py

@@ -1,42 +1,34 @@
-from typing import Any, io
+from typing import Any
 
 from IPython.core.display_functions import display
-from PIL import ImageShow
 from PIL.Image import Image
-
-original_show = ImageShow.show
+from PIL.ImageShow import UnixViewer
 
 
 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
+UnixViewer.show_file = show_file
 
-original_save = Image.save
+original_show = UnixViewer.show
 
 
-# To prevent circular save and display calls
-def __repr_image(self, image_format: str, **kwargs: Any) -> bytes | None:
-    """Helper function for iPython display hook.
+def show(self, image, **options):
+    display(image)
+    original_show(self, image, **options)
 
-    :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()
 
+UnixViewer.show = show
 
-Image._repr_image = __repr_image
+original_save = Image.save
 
 
 def save(image, fp, format=None, **options):
-    display(image)
+    if isinstance(fp, str):
+        display(image)
+
     original_save(image, fp, format, **options)