import base64
import struct
import tempfile
from contextlib import contextmanager
from io import BytesIO, StringIO
from typing import Any, BinaryIO, Dict, Generator, Optional, Tuple, Union
import matplotlib.pyplot as plt
from matplotlib.artist import Artist
from matplotlib.figure import Figure
from pybloqs.block.base import BaseBlock
from pybloqs.block.convenience import add_block_types
from pybloqs.html import append_to, parse
from pybloqs.static import Css, JScript
from pybloqs.util import cfg_to_css_string
try:
import plotly.offline as po
from plotly.graph_objs import Figure as PlotlyFigure
_PLOTLY_AVAILABLE = True
except ImportError:
_PLOTLY_AVAILABLE = False
try:
try:
# Bokeh >= 3.1
from bokeh.resources import Resources
JSResources = Resources
CSSResources = Resources
except ImportError:
try:
from bokeh.resources import CSSResources, JSResources
except ImportError:
from bokeh.templates import CSSResources, JSResources
try:
from bokeh.plotting.figure import Figure as BokehFigure
except ImportError:
from bokeh.plotting import figure as BokehFigure
from bokeh.embed.standalone import components
from bokeh.io import export_png
_BOKEH_AVAILABLE = True
except ImportError:
_BOKEH_AVAILABLE = False
_MIME_TYPES = {"png": "png", "svg": "svg+xml"}
_PLOT_FORMAT: str = "png"
_PLOT_MIME_TYPE = _MIME_TYPES[_PLOT_FORMAT]
_PLOT_DPI: int = 100
@contextmanager
def plot_format(plot_format: Optional[str] = None, dpi: Optional[int] = None) -> Generator[None, None, None]:
"""
Temporarily set the plot formatting settings
:param plot_format: The plot format (e.g 'png')
:type plot_format: str
:param dpi: The DPI of the plots
:type dpi: int
"""
old = get_plot_format()
set_plot_format(plot_format, dpi)
yield
set_plot_format(*old)
def get_plot_format() -> Tuple[str, int]:
"""
Get the current plot format parameters
:return: tuple of format and dpi
"""
return _PLOT_FORMAT, _PLOT_DPI
[docs]
class ImgBlock(BaseBlock):
def __init__(
self,
data: Union[BaseBlock, bytes],
mime_type: str = "png",
width: Optional[str] = None,
height: Optional[str] = None,
img_style: Optional[Dict] = None,
**kwargs,
) -> None:
"""
Create a block containing an image. The dimensions can be sniffed from GIF
and PNG data, for other formats, the `width` and `height` parameters can be
used to specify the size.
:param data: Either a block to convert to an image, or raw image data.
:mime_type: The mime type of the image (e.g. "png" or "svg+xml")
:param width: Image width override as string, e.g. "50px". If unspecified, the actual
image width is used.
:param width: Image height override as string, e.g. "50px". If unspecified, the actual
image height is used.
:param kwargs: Optional styling arguments. The `style` keyword argument has special
meaning in that it allows styling to be grouped as one argument.
It is also useful in case a styling parameter name clashes with a standard
block parameter.
"""
# Wrapping a block in an image will render it
if hasattr(data, "_write_block"):
img_file = data.save(fmt=mime_type)
with open(img_file, "rb") as f:
data = f.read()
self._img_data = base64.b64encode(data)
self._mime_type = mime_type
if width is None and height is None:
if mime_type.lower() == "png":
if struct.unpack("ccc", data[1:4]) != (b"P", b"N", b"G"):
raise ValueError("Image type is not png and does not match mime type")
x, y = struct.unpack(">ii", data[16:24])
elif mime_type.lower() == "gif":
x, y = struct.unpack("<HH", data[6:10])
else:
raise ValueError(f"Can't determine image dimensions for mime type {mime_type}")
width, height = (f"{x}px", f"{y}px")
if img_style is None:
img_styles = {}
else:
img_styles = img_style
if width is not None:
img_styles["width"] = width
if height is not None:
img_styles["height"] = height
self._img_styles = img_styles
super().__init__(**kwargs)
def _write_contents(self, container, *args, **kwargs) -> None:
src = StringIO()
mime = f"data:image/{self._mime_type};base64,"
src.write(mime)
src.write(self._img_data.decode())
img = append_to(container, "img", src=src.getvalue())
if len(self._img_styles) > 0:
img["style"] = cfg_to_css_string(self._img_styles)
[docs]
@staticmethod
def from_file(img_file: Union[str, BinaryIO], **kwargs) -> "ImgBlock":
"""
Load an image block from a file.
:param img_file: File path or file-like object.
:param kwargs: Arguments to pass to the `ImgBlock` constructor.
:return: ImgBlock instance.
"""
close_file = False
if isinstance(img_file, str):
img_file = open(img_file, "rb") # noqa: SIM115
close_file = True
try:
return ImgBlock(img_file.read(), **kwargs)
finally:
# Close the file in case it was opened in this function
if close_file:
img_file.close()
[docs]
class PlotBlock(ImgBlock):
def __init__(
self,
plot,
close_plot: bool = True,
bbox_inches="tight",
width: Optional[str] = None,
height: Optional[str] = None,
**kwargs,
) -> None:
"""
Create a block containing a matplotlib figure
:param plot: A matplotlib figure, axes or artist object.
:close_plot: Optional (default=True). Set to True to close the plot after it is
captured into an image and avoid lingering plot windows.
:bbox_inches: Optional bounding box parameter for 'figure.savefig'.
:param kwargs: Optional styling arguments. The `style` keyword argument has special
meaning in that it allows styling to be grouped as one argument.
It is also useful in case a styling parameter name clashes with a standard
block parameter.
"""
if not isinstance(plot, Artist):
raise ValueError("PlotBlock contents must be matplotlib Artist")
if isinstance(plot, Figure):
figure = plot
elif isinstance(plot, Artist):
figure = plot.get_figure()
else:
raise ValueError("Unexpected plot object type %s", type(plot))
img_data = BytesIO()
legends = []
for ax in figure.get_axes():
legend = ax.get_legend()
if legend is not None:
# Patch Legend get_window_extent since there seems to be a bug where
# it is passed an unexpected renderer instance.
_orig_get_window_extent = legend.get_window_extent
def _patched_get_window_extent(*_) -> Any:
return _orig_get_window_extent()
legend.get_window_extent = _patched_get_window_extent
legends.append(legend)
if len(figure.axes) == 0:
# empty plot, disable bbox_inches to that savefig still works
bbox_inches = None
figure.savefig(
img_data, dpi=_PLOT_DPI, format=_PLOT_FORMAT, bbox_extra_artists=legends, bbox_inches=bbox_inches
)
plt_width, plt_height = figure.get_size_inches()
width = width or f"{plt_width:0.3f}in"
height = height or f"{plt_height:0.3f}in"
if close_plot:
plt.close(figure)
super().__init__(img_data.getvalue(), _PLOT_MIME_TYPE, width=width, height=height, **kwargs)
def _to_static(self) -> BaseBlock:
# Convert to a basic image block in case we contain 'dynamic' svg content
return ImgBlock(self) if self._mime_type == "svg" else super()._to_static()
class PlotlyPlotBlock(BaseBlock):
def __init__(self, contents, plotly_kwargs: Optional[Dict[str, Any]] = None, static_kwargs=None, **kwargs) -> None:
"""
Writes out the content as raw text or HTML.
:param contents: Plotly graphics object figure.
:param plotly_kwargs: Kwargs that are passed to plotly plot function.
:param static_kwargs: Kwargs that are passed to plotly function write_image() for static output.
:param kwargs: Optional styling arguments. The `style` keyword argument has special
meaning in that it allows styling to be grouped as one argument.
It is also useful in case a styling parameter name clashes with a standard
block parameter.
"""
self.resource_deps = [JScript(script_string=po.offline.get_plotlyjs(), name="plotly")]
super().__init__(**kwargs)
if not isinstance(contents, PlotlyFigure):
raise ValueError("Expected plotly.graph_objs.graph_objs.Figure type but got %s", type(contents))
plotly_kwargs = plotly_kwargs or {}
self.static_kwargs = static_kwargs or {}
prefix = "<script>if (typeof require !== 'undefined') {var Plotly=require('plotly')}</script>"
self._fig = contents
self._contents = prefix + po.plot(contents, include_plotlyjs=False, output_type="div", **plotly_kwargs)
def _write_contents(self, container, *args, **kwargs) -> None:
container.append(parse(self._contents))
def _repr_html_(self) -> str:
return self.render_html()
def _to_static(self) -> ImgBlock:
# Create a static png image for use e.g. in an email body
with tempfile.NamedTemporaryFile(suffix=".png") as f:
self._fig.write_image(f.name, **self.static_kwargs)
static_block = ImgBlock.from_file(f)
return static_block
class BokehPlotBlock(BaseBlock):
def __init__(self, contents, static_kwargs=None, **kwargs) -> None:
"""
Writes out the content as raw text or HTML.
:param contents: Bokeh plotting figure.
:param static_kwargs: Kwargs that are passed to Bokeh function export_png() for static output.
:param kwargs: Optional styling arguments. The `style` keyword argument has special
meaning in that it allows styling to be grouped as one argument.
It is also useful in case a styling parameter name clashes with a standard
block parameter.
"""
self.resource_deps = [JScript(script_string=s, name="bokeh_js") for s in JSResources().js_raw]
self.resource_deps += [Css(css_string=s, name="bokeh_css") for s in CSSResources().css_raw]
super().__init__(**kwargs)
if not isinstance(contents, BokehFigure):
raise ValueError("Expected bokeh.plotting.figure.Figure type but got %s", type(contents))
self._fig = contents
self.static_kwargs = static_kwargs or {}
script, div = components(contents)
self._contents = script + div
def _write_contents(self, container, *args, **kwargs) -> None:
container.append(parse(self._contents))
def _to_static(self) -> ImgBlock:
# Create a static png image for use e.g. in an email body
with tempfile.NamedTemporaryFile(suffix=".png") as f:
export_png(self._fig, filename=f.name, **self.static_kwargs)
static_block = ImgBlock.from_file(f)
return static_block
add_block_types(Artist, PlotBlock)
# If Plotly or Bokeh are not installed skip registration
if _PLOTLY_AVAILABLE:
add_block_types(PlotlyFigure, PlotlyPlotBlock)
if _BOKEH_AVAILABLE:
add_block_types(BokehFigure, BokehPlotBlock)
