From dc3db3b1046197451f7adc0000a545370145a9a8 Mon Sep 17 00:00:00 2001
From: Kyle Goodrick <Kyle@Goodricks.net>
Date: Mon, 6 Jan 2025 14:21:35 -0700
Subject: [PATCH 1/2] Add polars LazyFrame extension for high-quality query
 plan graph output.

---
 marimo/__init__.py          |  3 ++
 marimo/_polars/lazyframe.py | 63 +++++++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+)
 create mode 100644 marimo/_polars/lazyframe.py

diff --git a/marimo/__init__.py b/marimo/__init__.py
index ded4cb59fd6..c62eeffc8c5 100644
--- a/marimo/__init__.py
+++ b/marimo/__init__.py
@@ -116,6 +116,9 @@
 from marimo._plugins.stateless.tabs import tabs
 from marimo._plugins.stateless.tree import tree
 from marimo._plugins.stateless.video import video
+from marimo._polars import (
+    lazyframe,  # noqa: F401  # Import is required for lazyframe registration
+)
 from marimo._runtime import output
 from marimo._runtime.capture import (
     capture_stderr,
diff --git a/marimo/_polars/lazyframe.py b/marimo/_polars/lazyframe.py
new file mode 100644
index 00000000000..cb1ad000545
--- /dev/null
+++ b/marimo/_polars/lazyframe.py
@@ -0,0 +1,63 @@
+import re
+from typing import TYPE_CHECKING, Any
+
+from marimo import mermaid
+from marimo._dependencies.dependencies import DependencyManager
+from marimo._output.hypertext import Html
+
+if TYPE_CHECKING:
+    import polars as pl
+
+if DependencyManager.polars.has():
+    import polars as pl
+
+    @pl.api.register_lazyframe_namespace("mo")
+    class Marimo:
+        def __init__(self, ldf: pl.LazyFrame):
+            self._ldf: pl.LazyFrame = ldf
+
+        def show_graph(self, **kwargs: Any) -> Html | str:
+            # We are specifying raw_output already, so we need to remove if it is passed
+            raw_output = kwargs.pop("raw_output", False)
+
+            dot = self._ldf.show_graph(raw_output=True, **kwargs)
+
+            if raw_output:
+                return dot
+
+            return self._dot_to_mermaid_html(dot)
+
+        # _dot_to_mermaid_html is a separate method from show_graph so that in the future
+        # polars can be updated to output mermaid directly when calling native show_graph
+        # inside a marimo environment.
+        # In order to do that we need a function that will not recursively call show_graph
+        @classmethod
+        def _dot_to_mermaid_html(self, dot: str) -> Html:
+            return mermaid(self._polars_dot_to_mermaid(dot))
+
+        @staticmethod
+        def _parse_node_label(label: str) -> str:
+            # replace escaped newlines
+            label = label.replace(r"\n", "\n")
+            # replace escaped quotes
+            label = label.replace('\\"', "#quot;")
+            return label
+
+        @classmethod
+        def _polars_dot_to_mermaid(cls, dot: str) -> str:
+            edge_regex = r"(?P<node1>\w+) -- (?P<node2>\w+)"
+            node_regex = r"(?P<node>\w+)(\s+)?\[label=\"(?P<label>.*)\"]"
+
+            nodes = [n for n in re.finditer(node_regex, dot)]
+            edges = [e for e in re.finditer(edge_regex, dot)]
+
+            return "\n".join(
+                [
+                    "graph TD",
+                    *[
+                        f'\t{n["node"]}["{cls._parse_node_label(n["label"])}"]'
+                        for n in nodes
+                    ],
+                    *[f'\t{e["node1"]} --- {e["node2"]}' for e in edges],
+                ]
+            )

From 28aa4122e6a03d2a54c34d6ff906ca52e9ca2095 Mon Sep 17 00:00:00 2001
From: Kyle Goodrick <Kyle@Goodricks.net>
Date: Mon, 6 Jan 2025 17:11:09 -0700
Subject: [PATCH 2/2] Add tests

---
 tests/_polars/test_polars.py | 69 ++++++++++++++++++++++++++++++++++++
 1 file changed, 69 insertions(+)
 create mode 100644 tests/_polars/test_polars.py

diff --git a/tests/_polars/test_polars.py b/tests/_polars/test_polars.py
new file mode 100644
index 00000000000..b1a115ad9b5
--- /dev/null
+++ b/tests/_polars/test_polars.py
@@ -0,0 +1,69 @@
+import pytest
+
+from marimo._dependencies.dependencies import DependencyManager
+from marimo._output.hypertext import Html
+
+HAS_DEPS = DependencyManager.polars.has()
+
+
+@pytest.fixture
+def simple_lf():
+    import polars as pl
+
+    return (
+        pl.LazyFrame(
+            {
+                "a": [1, 2, 3],
+                "b": [4, 5, 6],
+            }
+        )
+        .filter(pl.col("a") > 1)
+        .group_by("a")
+        .agg(pl.col("b").sum())
+    )
+
+
+@pytest.mark.skipif(not HAS_DEPS, reason="polars is required")
+def test_show_graph(simple_lf):
+    lf = simple_lf
+
+    assert type(lf.mo.show_graph()) is Html
+    assert type(lf.mo.show_graph(raw_output=True)) is str
+
+
+@pytest.mark.skipif(not HAS_DEPS, reason="polars is required")
+def test_dot_to_html(simple_lf):
+    lf = simple_lf
+
+    dot = lf.show_graph(raw_output=True)
+
+    assert type(lf.mo._dot_to_mermaid_html(dot)) is Html
+
+
+@pytest.mark.skipif(not HAS_DEPS, reason="polars is required")
+def test_parse_node_label(simple_lf):
+    lf = simple_lf
+
+    assert lf.mo._parse_node_label(r"\"") == "#quot;"
+    assert lf.mo._parse_node_label(r"\n") == "\n"
+    assert lf.mo._parse_node_label(r"\"\n\"") == "#quot;\n#quot;"
+
+
+@pytest.mark.skipif(not HAS_DEPS, reason="polars is required")
+def test_dot_to_mermaid(simple_lf):
+    lf = simple_lf
+    dot = lf.show_graph(raw_output=True)
+
+    mermaid_str = lf.mo._polars_dot_to_mermaid(dot)
+
+    assert type(mermaid_str) is str
+    assert mermaid_str == (
+        """graph TD
+	p2["TABLE
+π 2/2;
+σ [(col(#quot;a#quot;)) > (1)]"]
+	p1["AGG [col(#quot;b#quot;).sum()]
+BY
+[col(#quot;a#quot;)]"]
+	p1 --- p2"""
+    )