From 245db9aa8f50de5cd3d92f8fa4ab2091bb571c59 Mon Sep 17 00:00:00 2001
From: carmelom <carmelo.mordini@gmail.com>
Date: Tue, 15 Aug 2023 18:12:21 +0200
Subject: [PATCH] Add load_ options for Image component

---
 src/pydase/components/image.py | 62 +++++++++++++++++++++++-----------
 1 file changed, 42 insertions(+), 20 deletions(-)

diff --git a/src/pydase/components/image.py b/src/pydase/components/image.py
index 8d94ddf..13bb340 100644
--- a/src/pydase/components/image.py
+++ b/src/pydase/components/image.py
@@ -4,10 +4,18 @@ from pathlib import Path
 
 import PIL.Image
 from loguru import logger
+from urllib.request import urlopen
 
 from pydase.data_service.data_service import DataService
 
 
+class Figure:
+    """Mock class for matplotlib.Figure"""
+
+    def savefig(self, format="png"):
+        pass
+
+
 class Image(DataService):
     def __init__(
         self,
@@ -26,28 +34,42 @@ class Image(DataService):
 
     def load_from_path(self, path: Path | str) -> None:
         with PIL.Image.open(path) as image:
-            self._load_from_PIL_Image(image)
+            self._load_from_PIL(image)
 
-    def load_from_base64(self, value: bytes) -> None:
-        if isinstance(value, bytes):
-            # Decode the base64 string
-            image_data = base64.b64decode(value)
+    def load_from_matplotlib_figure(self, fig: Figure, format_: str = "png") -> None:
+        buffer = io.BytesIO()
+        fig.savefig(buffer, format=format_)
+        value_ = base64.b64encode(buffer.getvalue())
+        self._load_from_base64(value_, format_)
 
-            # Create a writable memory buffer for the image
-            image_buffer = io.BytesIO(image_data)
+    def load_from_url(self, url: str):
+        image = PIL.Image.open(urlopen(url))
+        self._load_from_PIL(image)
 
-            # Read the image from the buffer
-            image = PIL.Image.open(image_buffer)
-            self._load_from_PIL_Image(image)
+    def load_from_base64(self, value_: bytes, format_: str | None = None) -> None:
+        if format_ is None:
+            format_ = self._get_image_format_from_bytes(value_)
+        self._load_from_base64(value_, format_)
 
-    def _load_from_PIL_Image(self, image: PIL.Image.Image) -> None:
-        if isinstance(image, PIL.Image.Image):
-            if image.format is not None:
-                self._format = image.format
-                buffered = io.BytesIO()
-                image.save(buffered, format=self._format)
-                img_base64 = base64.b64encode(buffered.getvalue())
+    def _load_from_base64(self, value_: bytes, format_: str) -> None:
+        value = value_.decode("utf-8") if isinstance(value_, bytes) else value_
+        self._value = value
+        self._format = format_
 
-                self._value = img_base64.decode()
-            else:
-                logger.error("Image format is 'None'. Skipping...")
+    def _load_from_PIL(self, image: PIL.Image.Image) -> None:
+        if image.format is not None:
+            format_ = image.format
+            buffer = io.BytesIO()
+            image.save(buffer, format=format_)
+            value_ = base64.b64encode(buffer.getvalue())
+            self._load_from_base64(value_, format_)
+        else:
+            logger.error("Image format is 'None'. Skipping...")
+
+    def _get_image_format_from_bytes(self, value_: bytes):
+        image_data = base64.b64decode(value_)
+        # Create a writable memory buffer for the image
+        image_buffer = io.BytesIO(image_data)
+        # Read the image from the buffer
+        image = PIL.Image.open(image_buffer)
+        return image.format