Support PIL Images in pc.image (#1096)

pynecone/components/media/image.py

@@ -1,10 +1,12 @@
 """An image component."""
 from __future__ import annotations
 
-from typing import Optional, Set
+from typing import Any, Optional, Set
 
 from pynecone.components.component import Component
 from pynecone.components.libs.chakra import ChakraComponent
+from pynecone.components.tags import Tag
+from pynecone.utils import format, types
 from pynecone.vars import Var
 
 
@@ -38,7 +40,7 @@ class Image(ChakraComponent):
     loading: Var[str]
 
     # The image src attribute.
-    src: Var[str]
+    src: Var[Any]
 
     # The image srcset attribute.
     src_set: Var[str]
@@ -50,3 +52,11 @@ class Image(ChakraComponent):
             The event triggers.
         """
         return super().get_triggers() | {"on_error", "on_load"}
+
+    def _render(self) -> Tag:
+        # If the src is an image, convert it to a base64 string.
+        if types.is_image(type(self.src)):
+            self.src = Var.create(format.format_image_data(self.src))  # type: ignore
+
+        # Render the table.
+        return super()._render()

pynecone/utils/format.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import base64
+import io
 import json
 import os
 import re
@@ -359,6 +361,22 @@ def format_dataframe_values(value: Type) -> List[Any]:
     return format_data
 
 
+def format_image_data(value: Type) -> str:
+    """Format image data.
+
+    Args:
+        value: The value to format.
+
+    Returns:
+        Format data
+    """
+    buff = io.BytesIO()
+    value.save(buff, format="PNG")
+    image_bytes = buff.getvalue()
+    base64_image = base64.b64encode(image_bytes).decode("utf-8")
+    return f"data:image/png;base64,{base64_image}"
+
+
 def format_state(value: Any) -> Dict:
     """Recursively format values in the given state.
 
@@ -390,6 +408,10 @@ def format_state(value: Any) -> Dict:
             "data": format_dataframe_values(value),
         }
 
+    # Convert Image objects to base64.
+    if types.is_image(type(value)):
+        return format_image_data(value)  # type: ignore
+
     raise TypeError(
         "State vars must be primitive Python types, "
         "or subclasses of pc.Base. "

pynecone/utils/types.py

@@ -141,6 +141,20 @@ def is_dataframe(value: Type) -> bool:
     return value.__name__ == "DataFrame"
 
 
+def is_image(value: Type) -> bool:
+    """Check if the given value is a pillow image. By checking if the value subclasses PIL.
+
+    Args:
+        value: The value to check.
+
+    Returns:
+        Whether the value is a pillow image.
+    """
+    if is_generic_alias(value) or value == typing.Any:
+        return False
+    return "PIL" in value.__module__
+
+
 def is_figure(value: Type) -> bool:
     """Check if the given value is a figure.
 
@@ -162,7 +176,12 @@ def is_valid_var_type(var: Type) -> bool:
     Returns:
         Whether the value is a valid prop type.
     """
-    return _issubclass(var, StateVar) or is_dataframe(var) or is_figure(var)
+    return (
+        _issubclass(var, StateVar)
+        or is_dataframe(var)
+        or is_figure(var)
+        or is_image(var)
+    )
 
 
 def is_backend_variable(name: str) -> bool: