Source code for solidipes_core_plugin.viewers.pydantic

import streamlit as st
from IPython.display import Markdown as MarkdownIPython
from IPython.display import display
from solidipes.viewers import backends as viewer_backends
from solidipes.viewers.viewer import Viewer




class Pydantic(Viewer):
    """Viewer for Python modules defining Pydantic models."""

    def __init__(self, data=None):
        self.compatible_data_types = [str]
        self.module = None
        self.chart = None
        super().__init__(data)



    def add(self, data_container):
        """Append text to the viewer"""
        self.check_data_compatibility(data_container)
        self.module = data_container.module




    def show(self):
        if self.chart is None and self.module is not None:
            from pydantic_mermaid import MermaidGenerator, Relations

            generator = MermaidGenerator(self.module)
            self.chart = generator.generate_chart(relations=Relations.Both)

        if viewer_backends.current_backend == "jupyter notebook":
            display(MarkdownIPython(self.chart))

        elif viewer_backends.current_backend == "streamlit":
            import streamlit_mermaid as stmd

            try:
                stmd.st_mermaid(self.chart.replace("```mermaid", "").replace("```", ""))
            except Exception as e:
                st.error(f"Error rendering Mermaid chart: {e}")

        else:  # python
            print(self.chart)