fix test

jaytoday · Jul 14, 2023 · 8f62e66 · 8f62e66
1 parent b8c82e5
commit 8f62e66
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/examples/plot_sin_wave.py b/examples/plot_sin_wave.py
@@ -7,7 +7,20 @@ async def main():
         output = await session.generate_response(user_request)
 
         print(output.content)
-        print(output.files)
+
+        try:
+            from PIL import Image  # type: ignore
+        except ImportError:
+            print("Please install it with `pip install codeinterpreterapi[image_support]` to display images.")
+            exit(1)
+
+        from io import BytesIO
+        file = output.files[0]
+        img_io = BytesIO(file.content)
+        img = Image.open(img_io)
+
+        # Display the image
+        img.show()
 
 
 if __name__ == "__main__":