Source code for gofigr.backends.py3dmol

"""\
Copyright (c) 2023, Flagstaff Solutions, LLC
All rights reserved.

"""
import inspect
import os
import re
import sys
import tempfile

try:
    import py3Dmol
except ImportError:
    py3Dmol = None

try:
    from html2image import Html2Image
except ImportError:
    Html2Image = None

from gofigr.backends import GoFigrBackend, get_all_function_arguments




class WatermarkedView(py3Dmol.view):
    """Overrides py3Dmol.view to include a GoFigr watermark"""
    def __init__(self, wrapped_obj, extra_html):
        super().__init__()
        self.wrapped_obj = wrapped_obj
        self.extra_html = extra_html

    def _make_html(self):
        # pylint: disable=protected-access
        return self.wrapped_obj._make_html() + self.extra_html





class Py3DmolBackend(GoFigrBackend):
    """Plotly backend for GoFigr"""
    def __init__(self, *args, debug=False, image_size=(1920, 1080),
                 disable_logging=True, **kwargs):
        super().__init__(*args, **kwargs)
        self.image_size = image_size
        self.hti = Html2Image(size=image_size, disable_logging=disable_logging)
        self.debug = debug



    def get_backend_name(self):
        return "py3dmol"




    def is_compatible(self, fig):
        return isinstance(fig, py3Dmol.view)




    def is_interactive(self, fig):
        return True


    def _debug(self, text):
        if self.debug:
            print(text)

    def _infer_size(self, html):
        m = re.search(r'position: relative; width: (\d+)px; height: (\d+)px', html)
        if m is None:
            return self.image_size
        else:
            return int(m.group(1)), int(m.group(2))



    def find_figures(self, shell, data):
        self._debug("Looking for py3dmol figures...")
        self._debug(f"Available MIME types: {data.keys()}")

        frames = inspect.stack()
        # Make sure there's an actual figure being published, as opposed to Plotly initialization scripts
        if not any('3dmoljs' in x for x in data.keys()):
            self._debug("Could not locate 3dmoljs MIME type. Stopping search early.")
            return
        else:
            self._debug("3dmoljs MIME found")

        # Walk through the stack in *reverse* order (from top to bottom), to find the first call
        # in case display() was called recursively
        for idx, f in enumerate(reversed(frames)):
            self._debug(f"Frame {idx + 1}: function={f.function} file={f.filename}")
            if ("show" in f.function or "repr_html" in f.function) and "py3dmol" in f.filename.lower():
                self._debug("Found make_html call")
                for arg_idx, arg_value in enumerate(get_all_function_arguments(f)):
                    self._debug(f"Inspecting argument {arg_idx + 1}: "
                                f"{arg_value.__class__ if not arg_value is None else None}")
                    if self.is_compatible(arg_value):
                        self._debug("Argument is compatible with GoFigr")
                        yield arg_value
                    else:
                        self._debug("Argument not compatible.")

                break


    # pylint: disable=useless-return


    def get_default_figure(self, silent=False):
        if not silent:
            print("py3Dmol does not have a default figure. Please specify a figure to publish.", file=sys.stderr)

        return None




    def get_title(self, fig):
        title = getattr(fig, "title", None)
        if not isinstance(title, str):
            return None
        else:
            return title




    def figure_to_bytes(self, fig, fmt, params):
        if fmt == "png":
            with tempfile.NamedTemporaryFile(suffix=".png", dir=self.hti.output_path) as ntf:
                fig_html = self.figure_to_html(fig)
                width, height = self._infer_size(fig_html)
                self.hti.screenshot(html_str=fig_html, save_as=os.path.basename(ntf.name),
                                    size=(width, height))

                with open(ntf.name, 'rb') as f:
                    return f.read()

        raise ValueError(f"Py3Dmol does not support {fmt.upper()}")




    def figure_to_html(self, fig):
        return fig._make_html()  # pylint: disable=protected-access




    def add_interactive_watermark(self, fig, rev, watermark):
        watermark_html = f"<div><a href='{rev.revision_url}'>{rev.revision_url}</a></div>"
        return WatermarkedView(fig, watermark_html)




    def get_supported_image_formats(self):
        return ["png"]




    def close(self, fig):
        pass