feat(python): Add mermaid output to show_graph

py-polars/polars/_typing.py

@@ -187,6 +187,7 @@ def __arrow_c_stream__(self, requested_schema: object | None = None) -> object:
 TransferEncoding: TypeAlias = Literal["hex", "base64"]
 WindowMappingStrategy: TypeAlias = Literal["group_to_rows", "join", "explode"]
 ExplainFormat: TypeAlias = Literal["plain", "tree"]
+ShowGraphFormat: TypeAlias = Literal["dot", "mermaid"]
 
 # type signature for allowed frame init
 FrameInitTypes: TypeAlias = Union[

py-polars/polars/_utils/various.py

@@ -46,7 +46,7 @@
     from collections.abc import Iterator, MutableMapping, Reversible
 
     from polars import DataFrame, Expr
-    from polars._typing import PolarsDataType, SizeUnit
+    from polars._typing import PolarsDataType, ShowGraphFormat, SizeUnit
 
     if sys.version_info >= (3, 13):
         from typing import TypeIs
@@ -652,26 +652,68 @@ def re_escape(s: str) -> str:
     return re.sub(f"([{re_rust_metachars}])", r"\\\1", s)
 
 
+def dot_to_mermaid(dot: str) -> str:
+    """Not comprehensive, only handles components of the dot language used by polars."""
+    edge_regex = r"(?P<node1>\w+) -- (?P<node2>\w+)"
+    node_regex = r"(?P<node>\w+)(\s+)?\[label=\"(?P<label>.*)\"]"
+
+    nodes = re.finditer(node_regex, dot)
+    edges = re.finditer(edge_regex, dot)
+
+    mermaid_str = "\n".join(
+        [
+            "graph TD",
+            *[f'\t{n["node"]}["{n["label"]}"]' for n in nodes],
+            *[f'\t{e["node1"]} --- {e["node2"]}' for e in edges],
+        ]
+    )
+
+    # replace escaped newlines
+    mermaid_str = mermaid_str.replace(r"\n", "\n")
+
+    # replace escaped quotes
+    mermaid_str = mermaid_str.replace(r"\"", "#quot;")
+
+    return mermaid_str
+
+
 # Don't rename or move. This is used by polars cloud
 def display_dot_graph(
     *,
     dot: str,
     show: bool = True,
     output_path: str | Path | None = None,
     raw_output: bool = False,
+    raw_output_format: ShowGraphFormat = "dot",
     figsize: tuple[float, float] = (16.0, 12.0),
 ) -> str | None:
     if raw_output:
         # we do not show a graph, nor save a graph to disk
-        return dot
+        if raw_output_format == "dot":
+            return dot
+        elif raw_output_format == "mermaid":
+            return dot_to_mermaid(dot)
+        else:
+            msg = f"invalid raw output format: {raw_output_format}"
+            raise ValueError(msg)
 
     output_type = "svg" if _in_notebook() else "png"
 
     try:
         graph = subprocess.check_output(
             ["dot", "-Nshape=box", "-T" + output_type], input=f"{dot}".encode()
         )
+
+        graphviz_installed = True
     except (ImportError, FileNotFoundError):
+        graphviz_installed = False
+
+    exporting_file = output_path is not None
+    showing_outside_of_notebook = show and not _in_notebook()
+
+    if not graphviz_installed and (exporting_file or showing_outside_of_notebook):
+        # Graphviz must be installed to save graphs to disk or
+        # show them outside of a notebook
         msg = (
             "The graphviz `dot` binary should be on your PATH."
             "(If not installed you can download here: https://graphviz.org/download/)"
@@ -685,9 +727,14 @@ def display_dot_graph(
         return None
 
     if _in_notebook():
-        from IPython.display import SVG, display
-
-        return display(SVG(graph))
+        from IPython.display import SVG, Markdown, display
+
+        # if graphviz is installed, we can show the graph as an SVG
+        # otherwise we can show it as a mermaid diagram
+        if graphviz_installed:
+            return display(SVG(graph))
+        else:
+            return Markdown(f"```mermaid\n{dot_to_mermaid(dot)}\n```")
     else:
         import_optional(
             "matplotlib",

py-polars/polars/lazyframe/frame.py

@@ -118,6 +118,7 @@
         SchemaDefinition,
         SchemaDict,
         SerializationFormat,
+        ShowGraphFormat,
         StartBy,
         UniqueKeepStrategy,
     )
@@ -1131,6 +1132,7 @@ def show_graph(
         show: bool = True,
         output_path: str | Path | None = None,
         raw_output: bool = False,
+        raw_output_format: ShowGraphFormat = "dot",
         figsize: tuple[float, float] = (16.0, 12.0),
         type_coercion: bool = True,
         _type_check: bool = True,
@@ -1148,8 +1150,9 @@ def show_graph(
         """
         Show a plot of the query plan.
 
-        Note that Graphviz must be installed to render the visualization (if not
-        already present, you can download it here: `<https://graphviz.org/download>`_).
+        Note that Graphviz must be installed to export the visualization
+        or show it outside of a notebook
+        (if not already present, you can download it here: `<https://graphviz.org/download>`_).
 
         Parameters
         ----------
@@ -1161,6 +1164,8 @@ def show_graph(
             Write the figure to disk.
         raw_output
             Return dot syntax. This cannot be combined with `show` and/or `output_path`.
+        raw_output_format
+            The format of the raw output.
         figsize
             Passed to matplotlib if `show == True`.
         type_coercion
@@ -1221,6 +1226,7 @@ def show_graph(
             show=show,
             output_path=output_path,
             raw_output=raw_output,
+            raw_output_format=raw_output_format,
             figsize=figsize,
         )