diff --git a/pynecone/components/media/image.py b/pynecone/components/media/image.py index 455701429..4d62353e9 100644 --- a/pynecone/components/media/image.py +++ b/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() diff --git a/pynecone/utils/format.py b/pynecone/utils/format.py index a48be7203..f7c3507e3 100644 --- a/pynecone/utils/format.py +++ b/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. " diff --git a/pynecone/utils/types.py b/pynecone/utils/types.py index 4833ed611..cd75038a2 100644 --- a/pynecone/utils/types.py +++ b/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: